remove dfs, replace with replace_module

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Author: kylesayrs

examples/quantizing_moe/deepseek_r1_example.py

@@ -6,8 +6,13 @@
 from llmcompressor.transformers import oneshot
 
 # Select model and load it.
+
+# This script takes about 48 hours on 1xA100 to complete.
+# Future improvements will reduce this runtime (#1561, #1558).
+
 # For DeepSeek-R1, we require a full precision model in order to properly calibrate
 # `DeepSeek-R1-0528-BF16` is a DeepSeek-V3 FP8 model which has been converted to BF16
+
 model_id = "unsloth/DeepSeek-R1-0528-BF16"
 config = AutoConfig.from_pretrained(model_id)
 del config.quantization_config  # fp8 qconfig no longer appplies to bf16 model

src/llmcompressor/modeling/prepare.py

@@ -1,8 +1,7 @@
-import torch
+from compressed_tensors.utils import replace_module
 from transformers import PreTrainedModel
 
 from llmcompressor.modeling.deepseek_v3 import replace as replace_DeepseekV3MoE
-from llmcompressor.utils.module import module_bfs
 
 __all__ = ["prepare_for_calibration"]
 
@@ -12,11 +11,10 @@
 
 
 def prepare_for_calibration(model: PreTrainedModel) -> PreTrainedModel:
-    def replace(module: torch.nn.Module) -> torch.nn.Module:
+    for name, module in model.named_modules():
         cls_name = module.__class__.__name__
         if cls_name in replacements:
-            return replacements[cls_name](module)
-        else:
-            return module
+            new_module = replacements[cls_name](module)
+            replace_module(model, name, new_module)
 
-    return module_bfs(model, replace, progress=True)
+    return model