fix dist has_nan

Signed-off-by: Jennifer Chen <jennifchen@nvidia.com>

Author: jenchen13

examples/nemo_run/qat/README.md

@@ -61,12 +61,11 @@ To run the example locally, first clone the `TensorRT-Model-Optimizer` repositor
 Set up repo:
 
 - `git clone https://github.com/NVIDIA/TensorRT-Model-Optimizer.git`
-- `git clone https://github.com/NVIDIA-NeMo/NeMo.git`
 
 Run docker command (modify with your paths) and export the HuggingFace token:
 
 ```bash
-docker run -v  /home/user/:/home/user/ -v /home/user/NeMo:/opt/NeMo -v /home/user/TensorRT-Model-Optimizer/:/opt/TensorRT-Model-Optimizer/ --gpus all -it --shm-size 20g --rm nvcr.io/nvidia/nemo:25.09 bash
+docker run -v  /home/user/:/home/user/ -v /home/user/TensorRT-Model-Optimizer/:/opt/TensorRT-Model-Optimizer/ --gpus all -it --shm-size 20g --rm nvcr.io/nvidia/nemo:25.09 bash
 
 export HF_TOKEN=<your-token>
 ```

modelopt/torch/quantization/model_calib.py

@@ -26,7 +26,7 @@
 
 from modelopt.torch.opt.searcher import ForwardLoop
 from modelopt.torch.utils import print_rank_0
-from modelopt.torch.utils.distributed import ParallelState
+from modelopt.torch.utils.distributed import DistributedProcessGroup, ParallelState
 from modelopt.torch.utils.network import bind_forward_method, unpatch_forward_method
 
 from .conversion import create_and_replace_svdquant_linear_on_the_fly, set_quantizer_by_cfg_context
@@ -619,15 +619,11 @@ def sync_act_scale_across_dp(module, data_parallel_group):
             has_nan_local = torch.any(torch.isnan(module.awq_lite.act_scale)) or torch.any(
                 torch.isnan(module.awq_lite.weight_scale)
             )
-            has_nan = torch.tensor(int(has_nan_local), device=module.awq_lite.act_scale.device)
-            if module.parallel_state.data_parallel_group.is_initialized():
-                dist.all_reduce(
-                    has_nan,
-                    op=dist.ReduceOp.MAX,
-                    group=module.parallel_state.data_parallel_group.group,
-                )
+            has_nan = DistributedProcessGroup.get_dist_syncd_obj(
+                has_nan_local, module.parallel_state.data_parallel_group, lambda objs: any(objs)
+            )
 
-            if has_nan.item() > 0:
+            if has_nan:
                 module.awq_lite.is_enabled = False
             else:
                 sync_act_scale_across_dp(