NVIDIA-NeMo / NVIDIA-NeMo/RL

Ray data transfer bottleneck with long sequences/large batch size in distillation

Open
#1,877 0 comments 0 reactions 0 assignees View on GitHub
bug
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

When running distillation with long sequences, large batch size and topk settings (e.g., 32K response length, 512 batch size, 100 topk), Ray data transfer becomes a significant bottleneck, causing GPU underutilization.

**Affected Areas**

1. Training data dispatch (`policy_training/submit_training_futures`)
2. Teacher inference result return (`get_topk_logits` results)

**Observed Behavior**

With `max_total_sequence_length: 32768, train_global_batch_size: 512, topk_logits_k: 100`:

```other
Total step time: 945.93s

policy_training: 276.06s
├── submit_training_futures: 135.15s ← should be fast, but takes 135s

teacher_logprob_inference: 200.66s
├── (GPU mostly idle during result transfer)

generation: 192.38s (38.8%)
```

`submit_training_futures` should be a quick async dispatch, but it takes **135 seconds** due to large tensor serialization overhead.

**Steps/Code to reproduce bug**

The config for reproducing the above bug:

```yaml
defaults: distillation_math_megatron.yaml

distillation:
num_prompts_per_step: 512
max_num_steps: 1000
val_batch_size: 512
val_period: 20
val_at_start: false
topk_logits_k: 100

loss_fn:
kl_type: "reverse"

checkpointing:
checkpoint_dir: "checkpoints/distillation-8b-test"
save_period: 5

policy:
model_name: "Qwen/Qwen3-8B"
tokenizer:
name: "Qwen/Qwen3-8B"
train_global_batch_size: 512
generation_batch_size: 512
max_total_sequence_length: 32768
sequence_packing:
enabled: true
megatron_cfg:
tensor_model_parallel_size: 2
expert_tensor_parallel_size: 1
expert_model_parallel_size: 1
pipeline_model_parallel_size: 1
context_parallel_size: 4
sequence_parallel: true
generation:
vllm_cfg:
tensor_parallel_size: 4
expert_parallel_size: 1
make_sequence_length_divisible_by: 64
teacher:
model_name: "Qwen/Qwen3-8B"
tokenizer:
name: "Qwen/Qwen3-8B"
max_total_sequence_length: 32768
sequence_packing:
enabled: true
megatron_cfg:
tensor_model_parallel_size: 2
expert_tensor_parallel_size: 1
expert_model_parallel_size: 1
pipeline_model_parallel_size: 1
context_parallel_size: 4
sequence_parallel: true
make_sequence_length_divisible_by: 64

```

**Expected behavior**

Optimize the large data transfer.

**Additional context**

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.