Fix keep_in_fp32_modules not working for tied weights in load_and_quantize_model - #4043

Problem

MixedInt8LoaddedModelTest::test_fp32_8bit_conversion and Bnb4BitTestLoadedModel::test_fp32_4bit_conversion fail with:

AssertionError: assert torch.float16 == torch.float32

When quantizing an already-loaded model (not on meta device), modules specified in keep_in_fp32_modules (e.g. lm_head) are not correctly kept in fp32.

Root Cause

In BLOOM, lm_head.weight is tied to transformer.word_embeddings.weight (same Parameter object). model.named_parameters() deduplicates by parameter id and yields each parameter only once. Since transformer.word_embeddings.weight is encountered first in DFS traversal, it is the only one yielded — and the check "lm_head" in "transformer.word_embeddings.weight" is False, so the parameter stays in fp16. lm_head.weight is never iterated over.

Fix

Add a second pass using model.named_modules() to explicitly convert parameters of keep_in_fp32_modules to fp32. This is not affected by named_parameters() deduplication since we access the module's parameters directly.