minor update

Signed-off-by: Suguna Velury <[email protected]>

Author: sugunav14

examples/llm_ptq/multinode-ptq.py

@@ -50,8 +50,6 @@
 # Enable HuggingFace checkpointing
 mto.enable_huggingface_checkpointing()
 
-original_init_mp_dtypes = patch_fsdp_mp_dtypes()
-
 
 def parse_args():
     """Parse command line arguments."""
@@ -275,7 +273,7 @@ def export_model(
     export_dir.mkdir(parents=True, exist_ok=True)
 
     post_state_dict, hf_quant_config = _export_hf_checkpoint(
-        model, torch.bfloat16, is_fsdp2=True, accelerator=accelerator
+        model, torch.bfloat16, accelerator=accelerator
     )
 
     if accelerator.is_main_process:
@@ -384,9 +382,6 @@ def main(args):
         print(f"Model exported to {args.export_path}")
 
     print("Unpatching FSDP2 MP dtypes")
-    torch.distributed.fsdp._fully_shard._fsdp_param_group.FSDPParamGroup._init_mp_dtypes = (
-        original_init_mp_dtypes
-    )
 
 
 if __name__ == "__main__":

modelopt/torch/export/unified_export_hf.py

@@ -346,7 +346,6 @@ def _export_quantized_weight(
 def _export_hf_checkpoint(
     model: nn.Module,
     dtype: torch.dtype | None = None,
-    is_fsdp2: bool = False,
     accelerator: Accelerator | None = None,
 ) -> tuple[dict[str, Any], dict[str, Any]]:
     """Exports the torch model to the packed checkpoint with original HF naming.
@@ -356,6 +355,7 @@ def _export_hf_checkpoint(
     Args:
         model: the torch model.
         dtype: the weights data type to export the unquantized layers or the default model data type if None.
+        accelerator: the accelerator instance in case of distributed export setup.
 
     Returns:
         post_state_dict: Dict containing quantized weights
@@ -493,8 +493,7 @@ def _export_hf_checkpoint(
                     with fsdp2_aware_weight_update(model, sub_module):
                         _export_quantized_weight(sub_module, dtype, weight_name)
 
-    if is_fsdp2:
-        assert accelerator is not None, "Accelerator is required for FSDP2 export"
+    if accelerator is not None:
         # Gather state_dict from all ranks
         quantized_state_dict = accelerator.get_state_dict(model)
     else: