Ray data transfer bottleneck with long sequences/large batch size in distillation
- 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
Assessment
This issue has not been assessed yet.