NVIDIA-NeMo / NVIDIA-NeMo/RL

SFT training fails when using sequence_packing and TP=2 on 16K sequences

Open
#786 3 comments 0 reactions 1 assignee Claimed by @ahmadki View on GitHub
bug community-request external t-pytdensor x-google
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

**Describe the bug**

SFT training on 16K sequences fails when using sequence_packing and TP/SP is used. 8K sequences worked.

**Steps/Code to reproduce bug**

Inside `examples/configs/recipes/llm/sft-llama3.1-8b-instruct-1n8g-fsdp2tp2sp.v2.yaml`, replace the `sequence_packing` dict with the following:
```yaml
sequence_packing:
enabled: True
train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
algorithm: "modified_first_fit_decreasing"
sequence_length_round: 64
```
Run the following command:
```bash
uv run examples/run_sft.py --config examples/configs/recipes/llm/sft-llama3.1-8b-instruct-1n8g-fsdp2tp2sp.v2.yaml \
logger.log_dir=$LOCAL_LOG_DIR_IN_CONTAINER/ \
sft.val_global_batch_size=4 \
policy.train_global_batch_size=4 \
policy.max_total_sequence_length=16384 \
data.max_input_seq_length=16384 \
policy.dtensor_cfg.activation_checkpointing=true
```

**Expected behavior**
Encountered the following error:
```
[rank4]:[E729 18:02:00.338738893 ProcessGroupNCCL.cpp:1896] [PG ID 3 PG GUID 13(mesh_tp) Rank 0] Process group watchdog thread terminated with exception: CUDA error: an illegal memory access was encountered
```
Note that training with sequence length = 8192 works.

**Environment overview (please complete the following information)**

- Environment location: Docker
- Method of install: Building the Release docker container (https://github.com/NVIDIA-NeMo/RL/blob/main/docker/Dockerfile)
- If method of install is [Docker], provide `docker pull` & `docker run` commands used: N/A

**Environment details**

If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:
- OS version
- PyTorch version
- Python version

**Additional context**

Add any other context about the problem here.
Example: GPU model

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.