NVIDIA-NeMo / NVIDIA-NeMo/RL

Non-colocated with vllm tp8pp2 has abnormal token_mult_prob_error in FSDP2 backend

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

Description

vllm tp8pp2 could run well on:
1. colocated for both FSDP2 and Megatron backend.
2. non-colocated for Megatron backend.

vllm tp8pp2 could run well on non-colocated scenario on FSDP2 backend.

it's only abnormal when **vllm across node (model-parallel-size>8)** with **FSDP2 backend** with **non-colocated (nccl refit)**.

**Repro**
```bash
RUN_COMMAND="NRL_VLLM_ASYNC_TIMEOUT_SECONDS=14400 NRL_FORCE_REBUILD_VENVS=true uv run python examples/run_grpo_math.py \
--config examples/configs/grpo_math_8B.yaml \
grpo.max_num_steps=100 \
grpo.val_period=1000 \
policy.model_name=meta-llama/Llama-3.1-8B-Instruct \
policy.dtensor_cfg.enabled=true \
policy.dtensor_cfg.activation_checkpointing=true \
policy.dtensor_cfg.tensor_parallel_size=8 \
policy.generation.vllm_cfg.async_engine=true \
policy.generation.vllm_cfg.tensor_parallel_size=8 \
policy.generation.vllm_cfg.pipeline_parallel_size=2 \
policy.generation.colocated.enabled=false \
policy.generation.colocated.resources.num_nodes=4 \
policy.dynamic_batching.enabled=true \
policy.sequence_packing.enabled=false \
checkpointing.enabled=false \
logger.wandb_enabled=true \
logger.tensorboard_enabled=false \
logger.monitor_gpus=true \
logger.wandb.project=${PROJECT_NAME} \
logger.wandb.name=${EXP_NAME} \
cluster.num_nodes=6 \
cluster.gpus_per_node=8"
```

**Behavior**
purple: baseline
red: vllm tp8pp2 with `placement_group_strategy="SPREAD"` (same as basline)
blue: vllm tp8pp2 with `placement_group_strategy="PACK"`

Image

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.