SFT training fails when using sequence_packing and TP=2 on 16K sequences
- 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
Assessment
This issue has not been assessed yet.