Skip to content

fix(megatron): cast get_cp_local_num_tokens to torch.int for Megatron int32 token accumulator #290

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

Open
ZiyiTsang wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(megatron): cast get_cp_local_num_tokens to torch.int for Megatron int32 token accumulator #290

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
6 changes: 3 additions & 3 deletions relax/backends/megatron/cp_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -194,7 +194,7 @@ def get_cp_local_num_tokens(
"""
cp_size = dynamic_cp_size if dynamic_cp_size is not None else mpu.get_context_parallel_world_size()
if cp_size == 1:
return sum([torch.clamp_min(loss_mask.sum(), 1) for loss_mask in loss_masks])
return sum([torch.clamp_min(loss_mask.sum(), 1) for loss_mask in loss_masks]).to(torch.int)

# cp_size > 1: mirror the chunk slicing done in get_sum_of_sample_mean so the
# counted tokens exactly match the ones sum_of_token contributes on this rank.
Expand All @@ -221,8 +221,8 @@ def get_cp_local_num_tokens(

if total is None:
# No samples on this rank: mirror the empty-sum behaviour of cp_size == 1.
return sum([loss_mask.sum() for loss_mask in loss_masks])
return total
return sum([loss_mask.sum() for loss_mask in loss_masks]).to(torch.int)
return total.to(torch.int)


def all_gather_with_cp(
Expand Down
3 changes: 2 additions & 1 deletion tests/backends/megatron/test_rloo_policy_loss_dispatch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -167,7 +167,8 @@ def test_rloo_unequal_lengths_use_global_token_scalar_and_gradient_oracle(monkey
final_loss = token_sum_loss / normalizer

assert torch.allclose(token_sum_loss, expected_token_sum)
assert torch.allclose(normalizer, expected_num_tokens)
assert normalizer.dtype == torch.int
assert torch.allclose(normalizer.to(expected_num_tokens.dtype), expected_num_tokens)
assert torch.allclose(final_loss, expected_final_loss)

first_length = response_lengths[0]
Expand Down
Loading