git commit -m "Fix NCCL timeout issues in async checkpoint saving

- Add configure_nccl_for_checkpointing() to set appropriate timeouts
- Implement automatic fallback from async to sync on NCCL timeouts
- Configure NCCL_SOCKET_TIMEOUT and error handling for large models

Addresses NCCL timeout failures reported in #14576 when using
ckpt_async_save=True with large models like Qwen3 235B"

Signed-off-by: Abenezer <[email protected]>

Author: Abenez3r

nemo/utils/callbacks/dist_ckpt_io.py

@@ -58,6 +58,21 @@
         f" Exact error: {e}"
     )
 
+def configure_nccl_for_checkpointing():
+    """Configure NCCL settings for large model checkpoint saving."""
+    import os
+    
+    # Increase NCCL timeout for large model checkpointing
+    if not os.environ.get('NCCL_ASYNC_ERROR_HANDLING'):
+        os.environ['NCCL_ASYNC_ERROR_HANDLING'] = '1'
+    
+    # Set longer timeout (1 hour) for checkpoint operations
+    if not os.environ.get('NCCL_BLOCKING_WAIT'):
+        os.environ['NCCL_BLOCKING_WAIT'] = '1'
+    
+    # Increase socket timeout
+    if not os.environ.get('NCCL_SOCKET_TIMEOUT'):
+        os.environ['NCCL_SOCKET_TIMEOUT'] = '3600000'
 
 @contextmanager
 def _debug_time(name: str):
@@ -289,14 +304,34 @@ def save_checkpoint(
 
         rank = torch.distributed.get_rank()
         iteration = _get_iteration_from_checkpoint(checkpoint)
+        
+        # Configure NCCL for async saves to prevent timeouts
+        if self.async_save:
+            configure_nccl_for_checkpointing()
+        
         start_time = time()
-        async_save_request = dist_checkpointing.save(
-            sharded_state_dict=checkpoint,
-            checkpoint_dir=path,
-            sharded_strategy=self.save_sharded_strategy,
-            validate_access_integrity=validate_sharding_integrity,
-            async_sharded_save=self.async_save,
-        )
+        
+        try:
+            async_save_request = dist_checkpointing.save(
+                sharded_state_dict=checkpoint,
+                checkpoint_dir=path,
+                sharded_strategy=self.save_sharded_strategy,
+                validate_access_integrity=validate_sharding_integrity,
+                async_sharded_save=self.async_save,
+            )
+        except Exception as e:
+            if self.async_save and ("timeout" in str(e).lower() or "nccl" in str(e).lower()):
+                logging.warning(f"Async save failed with NCCL timeout ({e}), falling back to sync save")
+                async_save_request = dist_checkpointing.save(
+                    sharded_state_dict=checkpoint,
+                    checkpoint_dir=path,
+                    sharded_strategy=self.save_sharded_strategy,
+                    validate_access_integrity=validate_sharding_integrity,
+                    async_sharded_save=False,
+                )
+            else:
+                raise
+        
         end_time = time()
         log_parts = (
             "Global Checkpoint Save",