NVIDIA-NeMo / NVIDIA-NeMo/RL

[Bug] save_consolidated times out and truncates Gemma 4 31B safetensors with FSDP2/DTensor in v0.7.0

Open
#3,624 5 comments 0 reactions 2 assignees Claimed by @sharonyu-115 View on GitHub
bug
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.