From 4beb90e73694ea951a4cfcc2747f933edf7ef9ae Mon Sep 17 00:00:00 2001 From: Ziyi <1034337098@qq.com> Date: Mon, 24 Aug 2026 06:06:36 +0000 Subject: [PATCH] fix(megatron): cast get_cp_local_num_tokens to torch.int Both branches (cp_size == 1 and cp_size > 1) returned float32 (loss_mask.sum()), crashing every --calculate-per-token-loss training run at step 0: Megatron schedules.py accumulates `total_num_tokens += num_tokens` on an int tensor, so a float input raises "Result type Float can't be cast to the desired output type Int". Match the accumulator dtype, and compare the normalizer dtype-safely in the rloo dispatch test while pinning the int contract. --- relax/backends/megatron/cp_utils.py | 6 +++--- tests/backends/megatron/test_rloo_policy_loss_dispatch.py | 3 ++- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/relax/backends/megatron/cp_utils.py b/relax/backends/megatron/cp_utils.py index 98129f001..d088a3558 100644 --- a/relax/backends/megatron/cp_utils.py +++ b/relax/backends/megatron/cp_utils.py @@ -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. @@ -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( diff --git a/tests/backends/megatron/test_rloo_policy_loss_dispatch.py b/tests/backends/megatron/test_rloo_policy_loss_dispatch.py index a06d45c7b..fc1bfb962 100644 --- a/tests/backends/megatron/test_rloo_policy_loss_dispatch.py +++ b/tests/backends/megatron/test_rloo_policy_loss_dispatch.py @@ -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]