[Bug] save_consolidated times out and truncates Gemma 4 31B safetensors with FSDP2/DTensor in v0.7.0
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
## Describe the bug
In NeMo-RL v0.7.0, exporting a consolidated Hugging Face checkpoint for `google/gemma-4-31B-it` with AutoModel, FSDP2, and DTensor v2 fails when consolidation takes longer than the Gloo process group's 30-minute default timeout.
With `checkpointing.save_consolidated: true`, ranks synchronize through `dist.barrier()` inside `consolidate_safetensors_files_on_every_rank`. In NFS-backed storage, gathering and writing the 31B checkpoint takes approximately 35–45 minutes. The Gloo process group times out after 30 minutes while rank 0 is still writing.
The failed job leaves truncated `.safetensors` files. These files cannot subsequently be loaded by Hugging Face Transformers, NeMo AutoModel, or SGLang.
This is a regression. The same model, hardware, topology, launch configuration, storage, and checkpoint complete successfully in approximately 43 minutes using the earlier NeMo-RL commit `37cb7de1a`. The failure reproduces with NeMo-RL v0.7.0 and the intermediate `nemo-rl-arm64.130` build.
## Environment
- NeMo-RL failing version: `v0.7.0`
- Additional failing build: `nemo-rl-arm64.130`
- Last known working NeMo-RL commit: `37cb7de1a`
- Model: `google/gemma-4-31B-it`
- Model class: `Gemma4ForConditionalGeneration`
- Backend: NeMo AutoModel
- Distributed strategy: FSDP2 with DTensor v2
- Hardware: GB200
- Configurations tested:
- 32 GPUs: 8 workers × 4 GPUs
- 64 GPUs: 16 workers × 4 GPUs
- Parallelism: TP=1, CP=4
- CPU offload: enabled
- Distributed backend: `cuda:nccl,cpu:gloo`
- Model precision: BF16
- Storage: NFS-backed shared storage
- Python: 3.13
- Exact container image: `[ADD IMAGE AND TAG OR DIGEST]`
- PyTorch version: `[ADD VERSION]`
- Ray version: `[ADD VERSION]`
Checkpoint configuration:
```yaml
checkpointing:
save_optimizer: false
save_consolidated: true
```
## Steps to reproduce
1. Configure Gemma 4 31B training using the NeMo AutoModel FSDP2/DTensor v2 backend.
2. Use a distributed configuration equivalent to:
```yaml
tensor_parallel_size: 1
context_parallel_size: 4
cpu_offload: true
checkpointing:
save_optimizer: false
save_consolidated: true
```
3. Run on 32 or 64 GB200 GPUs using shared storage where consolidation takes longer than 30 minutes.
4. Allow training to save, or explicitly consolidate, a checkpoint.
5. Observe non-writing ranks enter `dist.barrier()` while rank 0 continues writing the Hugging Face safetensors.
6. After 30 minutes, observe the Gloo timeout:
```text
File ".../consolidate_hf_safetensors.py", line 855, in consolidate_safetensors_files_on_every_rank
dist.barrier()
File ".../torch/distributed/distributed_c10d.py", line 5037, in barrier
work.wait()
RuntimeError: [/pytorch/third_party/gloo/gloo/transport/tcp/unbound_buffer.cc:78]
Timed out waiting 1800000ms for recv operation to complete
```
7. Attempt to load the generated checkpoint using Transformers, AutoModel, or SGLang.
The load fails because one of the safetensors files is incomplete.
## Expected behavior
Consolidation should be allowed to run for the time required to export a large model.
The distributed timeout should be configurable and propagated to the process groups used by the checkpoint-consolidation path.
A failed consolidation must not leave truncated files that appear to be completed checkpoint artifacts.
## Actual behavior
The consolidation is terminated at the default 30-minute Gloo timeout.
The generated index declares a total model size of approximately 125 GB, but the first output shard is short by approximately 0.5–6.6 GB, depending on the run. The second shard completes successfully.
Loading the incomplete shard produces an error similar to:
```text
SafetensorError: Error while deserializing header:
incomplete metadata, file not fully covered
```
## Working workaround
Setting a longer timeout on every Ray worker before NeMo-RL initializes the distributed process group allows consolidation to complete:
```python
from datetime import timedelta
import torch.distributed as dist
_original_init_process_group = dist.init_process_group
def patched_init_process_group(*args, **kwargs):
if kwargs.get("timeout") is None:
kwargs["timeout"] = timedelta(hours=4)
return _original_init_process_group(*args, **kwargs)
dist.init_process_group = patched_init_process_group
```
With a four-hour timeout, the affected v0.7.0 consolidation completes successfully in approximately 43 minutes. This indicates that the checkpoint contents and storage path are otherwise functional.
## Suspected cause
NeMo-RL initializes the AutoModel process group without an explicit timeout:
```python
torch.distributed.init_process_group(backend=backend)
```
The CPU/Gloo group therefore uses PyTorch's 30-minute default. Large-model consolidation can legitimately exceed that timeout while other ranks wait at the post-write barrier.
The Python call appears unchanged between the working and failing builds. A bundled PyTorch, Gloo, Ray, or lower-level behavioral change may explain why the default timeout became observable in v0.7.0.
Contributor guide
Assessment
This issue has not been assessed yet.