Use torch get_rank() for broader coverage of use-cases

Signed-off-by: Jan Lasek <[email protected]>

nemo/export/quantize/quantizer.py

@@ -128,7 +128,7 @@ def _load_model(
 
         self._check_ddp_initialized(model)
 
-        if is_global_rank_zero():
+        if dist.get_rank() == 0:
             print(model)
 
         return model
@@ -183,7 +183,7 @@ def quantize(
 
         def forward_loop():
             for i, batch in enumerate(dataloader):
-                if is_global_rank_zero():
+                if dist.get_rank() == 0:
                     print(f"Calibrating batch {i}")
                 model.predict_step(batch, i)
 
@@ -212,7 +212,7 @@ def export(self, model, model_save: str):
                 export_tensorrt_llm_config=self.export_config.export_tensorrt_llm_config,
             )
             dist.barrier()  # Wait until all ranks complete export_model_config step
-            if is_global_rank_zero():
+            if dist.get_rank() == 0:
                 logging.info(f"Exporting quantized weights, model artifacts, and tokenizer config to {model_save}...")
                 save_artifacts(model, export_dir)
                 if save_qnemo: