Skip to content

fix: fix qwen3-235b deepseek-v3 h100 perf tests #2703

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
terrykong merged 8 commits into from
Jun 9, 2026
Merged

fix: fix qwen3-235b deepseek-v3 h100 perf tests #2703

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
51daa66
fix skip_reference_logprobs
yuki-97 Jun 5, 2026
6929534
use triton moe_backend
yuki-97 Jun 5, 2026
46336be
remove flaky check
yuki-97 Jun 5, 2026
cd840e7
fix rebase
yuki-97 Jun 6, 2026
68b68fa
update unit test
yuki-97 Jun 6, 2026
6439bb5
test(grpo): add setup auto-skip regression test and share mock_grpo_c…
yuki-97 Jun 9, 2026
aa6aad2
test: drop flaky step-10 token_mult_prob_error check from remaining h…
yuki-97 Jun 9, 2026
f0f9682
fix: pin moe_backend=triton in dapo-deepseek-v3 and document rationale
yuki-97 Jun 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions examples/configs/recipes/llm/performance/grpo-deepseek-v3-32n8g.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,8 @@ policy:
vllm_cfg:
tensor_parallel_size: 32
async_engine: true
vllm_kwargs:
moe_backend: triton
Comment thread
terrykong marked this conversation as resolved.
logger:
log_dir: logs/grpo-deepseek-v3-32n8g
wandb_enabled: true
Expand Down
2 changes: 2 additions & 0 deletions examples/configs/recipes/llm/performance/grpo-qwen3-235b-16n8g.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,8 @@ policy:
vllm_cfg:
tensor_parallel_size: 16
async_engine: true
vllm_kwargs:
moe_backend: triton
Comment thread
terrykong marked this conversation as resolved.
logger:
log_dir: logs/grpo-qwen3-235b-16n8g
wandb_enabled: true
Expand Down
73 changes: 41 additions & 32 deletions nemo_rl/algorithms/grpo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1634,7 +1634,6 @@ def grpo_train(
val_at_start = master_config.grpo["val_at_start"]
val_at_end = master_config.grpo["val_at_end"]
val_period = master_config.grpo["val_period"]
to_compute_kl = master_config.loss_fn.reference_policy_kl_penalty > 0
colocated_inference = master_config.policy["generation"]["colocated"]["enabled"]

# Initialize advantage estimator
Expand Down Expand Up @@ -2019,12 +2018,12 @@ def grpo_train(
"Computing prev_logprobs anyway for seq-level error masking."
)

if skip_prev_logprobs:
print(
"▶ Skipping prev_logprobs (force_on_policy_ratio=True)...",
flush=True,
)
else:
# Skip reference_policy_logprobs computation when skip_reference_policy_logprobs_calculation=True
skip_reference_logprobs = master_config.grpo.get(
"skip_reference_policy_logprobs_calculation"
)

if not (skip_prev_logprobs and skip_reference_logprobs):
print("▶ Preparing for logprob inference...", flush=True)
with timer.time("logprob_inference_prep"):
policy.prepare_for_lp_inference()
Expand All @@ -2041,25 +2040,32 @@ def grpo_train(
**extra_multimodal_data,
}
)

if not skip_prev_logprobs:
train_data["prev_logprobs"] = policy.get_logprobs(
logprob_data, timer=timer
)["logprobs"]
else:
print(
"▶ Skipping prev_logprobs (force_on_policy_ratio=True)...",
flush=True,
)
train_data["prev_logprobs"] = torch.zeros_like(
train_data["generation_logprobs"]
)

if to_compute_kl and not master_config.grpo.get(
"skip_reference_policy_logprobs_calculation"
):
if not skip_reference_logprobs:
train_data["reference_policy_logprobs"] = (
policy.get_reference_policy_logprobs(
logprob_data,
timer=timer,
)["reference_logprobs"]
)
else:
print(
"▶ Skipping reference_logprobs (skip_reference_policy_logprobs_calculation=True)...",
flush=True,
)
train_data["reference_policy_logprobs"] = torch.zeros_like(
train_data["prev_logprobs"]
)
Expand Down Expand Up @@ -2800,7 +2806,6 @@ def async_grpo_train(
val_period = master_config.grpo["val_period"]
val_at_start = master_config.grpo["val_at_start"]
val_at_end = master_config.grpo["val_at_end"]
to_compute_kl = master_config.loss_fn.reference_policy_kl_penalty > 0
colocated_inference = master_config.policy["generation"]["colocated"]["enabled"]

# Initialize advantage estimator
Expand Down Expand Up @@ -3172,37 +3177,41 @@ def async_grpo_train(
"Computing prev_logprobs anyway for seq-level error masking."
)

if skip_prev_logprobs:
print(
"▶ Skipping prev_logprobs (force_on_policy_ratio=True)...",
flush=True,
)
fprop_logprobs = torch.zeros_like(train_data["generation_logprobs"])
else:
print("▶ Preparing for logprob inference...")
# Skip reference_policy_logprobs computation when skip_reference_policy_logprobs_calculation=True
skip_reference_logprobs = master_config.grpo.get(
"skip_reference_policy_logprobs_calculation"
)

if not (skip_prev_logprobs and skip_reference_logprobs):
print("▶ Preparing for logprob inference...", flush=True)
with timer.time("logprob_inference_prep"):
policy.prepare_for_lp_inference()

print("▶ Computing logprobs...", flush=True)
with timer.time("policy_and_reference_logprobs"):
if not skip_prev_logprobs:
fprop_logprobs = policy.get_logprobs(
train_data,
timer=timer,
train_data["prev_logprobs"] = policy.get_logprobs(
train_data, timer=timer
)["logprobs"]
train_data["prev_logprobs"] = fprop_logprobs
else:
train_data["prev_logprobs"] = torch.zeros_like(
train_data["generation_logprobs"]
)

if to_compute_kl and not master_config.grpo.get(
"skip_reference_policy_logprobs_calculation"
):
reference_logprobs = policy.get_reference_policy_logprobs(
train_data,
timer=timer,
)["reference_logprobs"]
train_data["reference_policy_logprobs"] = reference_logprobs
if not skip_reference_logprobs:
train_data["reference_policy_logprobs"] = (
policy.get_reference_policy_logprobs(
train_data,
timer=timer,
)["reference_logprobs"]
)
else:
print(
"▶ Skipping reference_logprobs (skip_reference_policy_logprobs_calculation=True)...",
flush=True,
)
train_data["reference_policy_logprobs"] = torch.zeros_like(
fprop_logprobs
train_data["prev_logprobs"]
)

# Seq-level logprob error metrics/masking require real prev_logprobs
Expand Down
3 changes: 1 addition & 2 deletions tests/test_suites/llm/performance/grpo-qwen3-235b-16n8g.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,8 +36,7 @@ uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
# Only run metrics if the target step is reached
if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then
uv run tests/check_metrics.py $JSON_METRICS \
'median(data["train/token_mult_prob_error"]) < 1.1' \
'data["train/token_mult_prob_error"]["10"] < 1.1'
'median(data["train/token_mult_prob_error"]) < 1.1'
Comment thread
terrykong marked this conversation as resolved.

# Clean up checkpoint directory after successful run to save space.
rm -rf "$CKPT_DIR"
Expand Down
25 changes: 7 additions & 18 deletions tests/unit/algorithms/test_grpo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1811,28 +1811,17 @@ def fake_batched_message_log_to_flat_message(*_args, **_kwargs):


@pytest.mark.parametrize("train_func", [grpo_train, async_grpo_train])
@pytest.mark.parametrize(
"skip_reference_policy_logprobs_calculation",
[False, True],
ids=["implicit_kl_zero", "explicit_skip"],
)
def test_grpo_train_skips_reference_policy_logprobs_when_kl_disabled(
mock_grpo_components, train_func, skip_reference_policy_logprobs_calculation
):
def test_grpo_train_skips_reference_policy_logprobs(mock_grpo_components, train_func):
Comment thread
terrykong marked this conversation as resolved.
"""Regression test for issue #1968 (Bug 1) and PRs #2174 / #2178.

When ``loss_fn.reference_policy_kl_penalty=0``, both ``grpo_train`` and
``async_grpo_train`` MUST NOT call ``policy.get_reference_policy_logprobs``,
even if the caller did not pre-set
``grpo.skip_reference_policy_logprobs_calculation``. Without the skip guards
in ``grpo.py``, training would crash inside ``use_reference_model()`` because
the reference model state was never loaded.
When skip_reference_policy_logprobs_calculation=True, both grpo_train
and async_grpo_train MUST NOT call policy.get_reference_policy_logprobs.
Without the skip guards in grpo.py, training would crash inside
use_reference_model() because the reference model state was never loaded.
"""
master_config = mock_grpo_components["master_config"]
master_config.loss_fn.reference_policy_kl_penalty = 0
master_config.grpo["skip_reference_policy_logprobs_calculation"] = (
skip_reference_policy_logprobs_calculation
)
master_config.grpo["skip_reference_policy_logprobs_calculation"] = True
master_config.grpo["max_num_steps"] = 1
master_config.grpo["max_num_epochs"] = 1
master_config.grpo["val_period"] = 0
Expand Down Expand Up @@ -1903,7 +1892,7 @@ def test_grpo_train_skips_reference_policy_logprobs_when_kl_disabled(

assert not policy.get_reference_policy_logprobs.called, (
"policy.get_reference_policy_logprobs was called even though "
"loss_fn.reference_policy_kl_penalty=0. "
"skip_reference_policy_logprobs_calculation=True. "
"This indicates a regression of issue #1968 / PRs #2174, #2178."
)

Expand Down
Loading