Skip to content

fix: Fixes to make Megatron backend match dtensor #1389

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 19 commits into from
Oct 31, 2025
Merged

fix: Fixes to make Megatron backend match dtensor #1389

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
19 commits
Select commit Hold shift + click to select a range
ce6fd8a
enable megatron calculate_per_token_loss
ashors1 Oct 16, 2025
bdcee3c
remove unneeded import
ashors1 Oct 17, 2025
408900a
TP bug fix and add WIP unit test
ashors1 Oct 23, 2025
d8348d1
fixes
ashors1 Oct 24, 2025
9d0c410
lint
ashors1 Oct 24, 2025
72a5c73
add tensor parallel attributes check
ashors1 Oct 28, 2025
dcdd016
add convergence test
ashors1 Oct 28, 2025
9f99d53
address comment
ashors1 Oct 28, 2025
5cfad4a
Merge branch 'main' of github.com:NVIDIA-NeMo/RL into ashors/fix-mcor…
ashors1 Oct 28, 2025
3007940
improve unit test
ashors1 Oct 28, 2025
ba86f3a
minimize config
ashors1 Oct 28, 2025
37d57d6
fix permissions
ashors1 Oct 28, 2025
7a65797
fix convergence test
ashors1 Oct 28, 2025
198027a
minimize config
ashors1 Oct 29, 2025
64594d0
fix unit test
ashors1 Oct 29, 2025
f2cb2ba
remove more occurrences of average_in_collective
ashors1 Oct 30, 2025
cc8b6d7
link to GH issue
ashors1 Oct 30, 2025
d88f478
Merge branch 'main' of github.com:NVIDIA-NeMo/RL into ashors/fix-mcor…
ashors1 Oct 31, 2025
f12e113
fix test failure
ashors1 Oct 31, 2025
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
1 change: 0 additions & 1 deletion examples/configs/dpo.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -153,7 +153,6 @@ policy:
grad_reduce_in_fp32: false
overlap_grad_reduce: true
overlap_param_gather: true
average_in_collective: true
data_parallel_sharding_strategy: "optim_grads_params"

data:
Expand Down
1 change: 0 additions & 1 deletion examples/configs/grpo_math_1B.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,6 @@ policy:
grad_reduce_in_fp32: false
overlap_grad_reduce: true
overlap_param_gather: true
average_in_collective: true
use_custom_fsdp: false
data_parallel_sharding_strategy: "optim_grads_params"

Expand Down
1 change: 0 additions & 1 deletion examples/configs/grpo_math_1B_megatron.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,7 +133,6 @@ policy:
grad_reduce_in_fp32: false
overlap_grad_reduce: true
overlap_param_gather: true
average_in_collective: true
use_custom_fsdp: false
data_parallel_sharding_strategy: "optim_grads_params"

Expand Down
1 change: 0 additions & 1 deletion examples/configs/rm.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,6 @@ policy:
grad_reduce_in_fp32: false
overlap_grad_reduce: true
overlap_param_gather: false
average_in_collective: true
data_parallel_sharding_strategy: "optim_grads_params"


Expand Down
1 change: 0 additions & 1 deletion examples/configs/sft.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -134,7 +134,6 @@ policy:
grad_reduce_in_fp32: false
overlap_grad_reduce: true
overlap_param_gather: true
average_in_collective: true
data_parallel_sharding_strategy: "optim_grads_params"
use_custom_fsdp: false

Expand Down
1 change: 0 additions & 1 deletion examples/configs/sft_openmathinstruct2_megatron.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,6 @@ policy:
activation_checkpointing: false
context_parallel_size: 1
distributed_data_parallel_config:
average_in_collective: true
data_parallel_sharding_strategy: optim_grads_params
grad_reduce_in_fp32: true
overlap_grad_reduce: true
Expand Down
1 change: 0 additions & 1 deletion examples/configs/vlm_grpo_3B.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,6 @@ policy:
grad_reduce_in_fp32: false
overlap_grad_reduce: true
overlap_param_gather: true
average_in_collective: true
use_custom_fsdp: false
data_parallel_sharding_strategy: "optim_grads_params"

Expand Down
1 change: 0 additions & 1 deletion examples/configs/vlm_grpo_3B_megatron.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,6 @@ policy:
grad_reduce_in_fp32: false
overlap_grad_reduce: false
overlap_param_gather: true
average_in_collective: true
use_custom_fsdp: false
data_parallel_sharding_strategy: optim_grads_params
data:
Expand Down
7 changes: 4 additions & 3 deletions nemo_rl/models/policy/megatron_policy_worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -658,6 +658,9 @@ def __init__(
"https://github.com/NVIDIA-NeMo/RL/blob/bccbc377705a81a1f4b3c31ad9767bcc15f735a8/nemo_rl/algorithms/sft.py#L175-L179."
)

## TODO: make sure this works with sequence-level losses as well
Comment thread
ashors1 marked this conversation as resolved.
Outdated
model_cfg.calculate_per_token_loss = True

Comment thread
ashors1 marked this conversation as resolved.
self.megatron_cfg = ConfigContainer(
model=model_cfg,
checkpoint=checkpoint_config,
Expand All @@ -683,9 +686,7 @@ def __init__(
overlap_param_gather=self.cfg["megatron_cfg"][
"distributed_data_parallel_config"
]["overlap_param_gather"],
average_in_collective=self.cfg["megatron_cfg"][
"distributed_data_parallel_config"
]["average_in_collective"],
average_in_collective=False, # average in collective is not supported with per-token loss
Comment thread
terrykong marked this conversation as resolved.
Outdated
use_distributed_optimizer=self.cfg["megatron_cfg"]["optimizer"][
"use_distributed_optimizer"
],
Expand Down
Loading