From e9d4279324ed9069e8b649567a9a524779ae3502 Mon Sep 17 00:00:00 2001 From: luozijian Date: Fri, 11 Sep 2026 00:03:55 +0800 Subject: [PATCH 1/2] [Bugfix][Training] Preserve DSpark rotary embedding dtype Keep rotary tables in floating point instead of casting them to long position IDs, then match the query dtype at the attention boundary. Add focused CPU regression coverage and premerge registration. AI-assisted: developed and reviewed with ChatGPT; submitter must add DCO sign-off before publication. Signed-off-by: luozijian --- .buildkite/pipeline.yml | 1 + tests/test_dspark_rotary_embedding.py | 52 +++++++++++++++++++ .../megatron_utils/dspark/attention.py | 4 +- 3 files changed, 56 insertions(+), 1 deletion(-) create mode 100644 tests/test_dspark_rotary_embedding.py diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index dbc8b225..3ac5f59b 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -92,6 +92,7 @@ steps: python tests/test_expert_routing.py python tests/test_reloadable_process_group_memory_check.py python tests/test_ppo_logprob_entropy.py + python tests/test_dspark_rotary_embedding.py python tests/utils/test_hf_checkpoint_saver.py ' diff --git a/tests/test_dspark_rotary_embedding.py b/tests/test_dspark_rotary_embedding.py new file mode 100644 index 00000000..68719745 --- /dev/null +++ b/tests/test_dspark_rotary_embedding.py @@ -0,0 +1,52 @@ +import math + +import torch + +from vime.backends.megatron_utils.dspark.attention import DSparkParallelAttention, DSparkRotaryEmbedding + + +NUM_GPUS = 0 + + +def test_dspark_rotary_embedding_keeps_long_position_ids_out_of_output_dtype(): + rotary = DSparkRotaryEmbedding(head_dim=8, rotary_base=10000.0) + position_ids = torch.tensor([[0, 1, 17, 128]], dtype=torch.long) + + cos, sin = rotary(position_ids) + + inv_freq = 1.0 / (10000.0 ** (torch.arange(0, 8, 2, dtype=torch.float32) / 8)) + freqs = torch.einsum("i,bj->bji", inv_freq, position_ids.float()) + emb = torch.cat([freqs, freqs], dim=-1) + expected_cos = emb.cos().unsqueeze(1) + expected_sin = emb.sin().unsqueeze(1) + + assert cos.dtype == torch.float32 + assert sin.dtype == torch.float32 + torch.testing.assert_close(cos, expected_cos, rtol=0, atol=0) + torch.testing.assert_close(sin, expected_sin, rtol=0, atol=0) + assert torch.count_nonzero(sin).item() > 0 + assert math.isclose(float(cos.abs().max()), 1.0, rel_tol=0, abs_tol=1e-7) + + +def test_dspark_attention_casts_rotary_values_to_query_dtype(): + attention = DSparkParallelAttention( + hidden_size=256, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=64, + ).to(dtype=torch.bfloat16) + hidden_states = torch.randn(2, 3, 256, dtype=torch.bfloat16) + target_hidden_states = torch.randn(2, 5, 256, dtype=torch.bfloat16) + position_ids = torch.arange(8, dtype=torch.long).unsqueeze(0).expand(2, -1) + attention_mask = torch.ones((2, 1, 3, 8), dtype=torch.bool) + + output = attention( + hidden_states=hidden_states, + target_hidden_states=target_hidden_states, + position_ids=position_ids, + attention_mask=attention_mask, + ) + + assert output.dtype == torch.bfloat16 + assert output.shape == hidden_states.shape + assert torch.isfinite(output).all() diff --git a/vime/backends/megatron_utils/dspark/attention.py b/vime/backends/megatron_utils/dspark/attention.py index 2378260b..b14043b0 100644 --- a/vime/backends/megatron_utils/dspark/attention.py +++ b/vime/backends/megatron_utils/dspark/attention.py @@ -86,7 +86,7 @@ def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tenso # Add head dim for broadcasting: [bsz, 1, seq_len, head_dim] cos = cos.unsqueeze(1) sin = sin.unsqueeze(1) - return cos.to(position_ids.dtype), sin.to(position_ids.dtype) + return cos, sin class DSparkParallelAttention(nn.Module): @@ -169,6 +169,8 @@ def forward( # Apply rotary embeddings # position_ids: [bsz, ctx_len + q_len] cos, sin = self.rotary_emb(position_ids) + cos = cos.to(dtype=q.dtype) + sin = sin.to(dtype=q.dtype) # cos/sin: [1, 1, seq_len, head_dim] but we need to match k's shape # k shape: [bsz, kv_heads, kv_len, head_dim] # cos shape: [1, 1, kv_len, head_dim] -> broadcast over bsz and heads From cfac5c2bacf6dcfe72c96b807a06c2da3d622dbe Mon Sep 17 00:00:00 2001 From: luozijian Date: Fri, 11 Sep 2026 09:27:44 +0800 Subject: [PATCH 2/2] [Bugfix][Training] Preserve DSpark FP32 phases across dtype conversion Recompute inverse frequencies on the positions device and avoid autocast downcasting of the outer product. Cover parent casts, long positions, attention gradients, and CPU/CUDA autocast. Execute the CPU suite via pytest in Buildkite rather than merely importing its definitions. Assisted-by: ChatGPT Signed-off-by: luozijian --- .buildkite/pipeline.yml | 2 +- tests/test_dspark_rotary_embedding.py | 80 ++++++++++++++++++- .../megatron_utils/dspark/attention.py | 17 ++-- 3 files changed, 88 insertions(+), 11 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 3ac5f59b..fdc72411 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -92,7 +92,7 @@ steps: python tests/test_expert_routing.py python tests/test_reloadable_process_group_memory_check.py python tests/test_ppo_logprob_entropy.py - python tests/test_dspark_rotary_embedding.py + python -m pytest tests/test_dspark_rotary_embedding.py -q python tests/utils/test_hf_checkpoint_saver.py ' diff --git a/tests/test_dspark_rotary_embedding.py b/tests/test_dspark_rotary_embedding.py index 68719745..cbb55222 100644 --- a/tests/test_dspark_rotary_embedding.py +++ b/tests/test_dspark_rotary_embedding.py @@ -1,5 +1,7 @@ +import copy import math +import pytest import torch from vime.backends.megatron_utils.dspark.attention import DSparkParallelAttention, DSparkRotaryEmbedding @@ -8,6 +10,15 @@ NUM_GPUS = 0 +def _reference_rotary(position_ids, head_dim, rotary_base): + inv_freq = 1.0 / ( + rotary_base ** (torch.arange(0, head_dim, 2, device=position_ids.device, dtype=torch.float32) / head_dim) + ) + angles = position_ids.float().unsqueeze(-1) * inv_freq + angles = torch.cat([angles, angles], dim=-1).unsqueeze(1) + return angles.cos(), angles.sin() + + def test_dspark_rotary_embedding_keeps_long_position_ids_out_of_output_dtype(): rotary = DSparkRotaryEmbedding(head_dim=8, rotary_base=10000.0) position_ids = torch.tensor([[0, 1, 17, 128]], dtype=torch.long) @@ -22,12 +33,77 @@ def test_dspark_rotary_embedding_keeps_long_position_ids_out_of_output_dtype(): assert cos.dtype == torch.float32 assert sin.dtype == torch.float32 - torch.testing.assert_close(cos, expected_cos, rtol=0, atol=0) - torch.testing.assert_close(sin, expected_sin, rtol=0, atol=0) + torch.testing.assert_close(cos, expected_cos, rtol=1e-5, atol=1e-6) + torch.testing.assert_close(sin, expected_sin, rtol=1e-5, atol=1e-6) assert torch.count_nonzero(sin).item() > 0 assert math.isclose(float(cos.abs().max()), 1.0, rel_tol=0, abs_tol=1e-7) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("head_dim,rotary_base", [(64, 10000.0), (128, 500000.0)]) +@pytest.mark.parametrize("roundtrip", [False, True]) +def test_rotary_remains_fp32_after_parent_dtype_conversion(dtype, head_dim, rotary_base, roundtrip): + parent = torch.nn.Module() + parent.rotary = DSparkRotaryEmbedding(head_dim, rotary_base) + parent.to(dtype=dtype) + if roundtrip: + parent.float() # Widening an already-rounded buffer cannot recover it. + position_ids = torch.tensor([[0, 1, 128, 1024, 8192, 32768], [32769, 4096, 17, 2, 1, 0]], dtype=torch.long) + + actual = parent.rotary(position_ids) + expected = _reference_rotary(position_ids, head_dim, rotary_base) + + assert not parent.state_dict() # No persistent rotary checkpoint keys. + for value, reference in zip(actual, expected, strict=True): + assert value.dtype == torch.float32 + torch.testing.assert_close(value, reference, rtol=1e-5, atol=1e-6) + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("autocast_dtype", [torch.float16, torch.bfloat16]) +def test_rotary_autocast_preserves_fp32_phases(device, autocast_dtype): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA is required for the CUDA autocast regression") + rotary = DSparkRotaryEmbedding(head_dim=64).to(device=device) + positions = torch.tensor([[0, 128, 8192, 32768]], dtype=torch.long, device=device) + expected = _reference_rotary(positions, 64, 10000.0) + with torch.autocast(device_type=device, dtype=autocast_dtype): + actual = rotary(positions) + for value, reference in zip(actual, expected, strict=True): + assert value.dtype == torch.float32 + torch.testing.assert_close(value, reference, rtol=1e-5, atol=1e-6) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_low_precision_attention_matches_fresh_fp32_rotary_forward_and_backward(dtype): + with torch.random.fork_rng(devices=[]): + torch.manual_seed(31) + attention = DSparkParallelAttention( + hidden_size=256, num_attention_heads=4, num_key_value_heads=2, head_dim=64 + ).to(dtype=dtype) + reference = copy.deepcopy(attention) + reference.rotary_emb = DSparkRotaryEmbedding(head_dim=64) + hidden = torch.randn(2, 3, 256, dtype=dtype, requires_grad=True) + context = torch.randn(2, 5, 256, dtype=dtype, requires_grad=True) + reference_hidden = hidden.detach().clone().requires_grad_() + reference_context = context.detach().clone().requires_grad_() + positions = torch.tensor([[0, 17, 128, 4096, 8192, 16384, 32768, 32769]], dtype=torch.long).expand(2, -1) + mask = torch.ones((2, 1, 3, 8), dtype=torch.bool) + + actual = attention(hidden, context, positions, mask) + expected = reference(reference_hidden, reference_context, positions, mask) + assert actual.dtype == dtype + torch.testing.assert_close(actual.float(), expected.float(), rtol=1e-5, atol=1e-6) + actual.float().square().mean().backward() + expected.float().square().mean().backward() + for actual_input, expected_input in ((hidden, reference_hidden), (context, reference_context)): + assert actual_input.grad is not None and expected_input.grad is not None + torch.testing.assert_close(actual_input.grad.float(), expected_input.grad.float(), rtol=1e-5, atol=1e-6) + for actual_param, expected_param in zip(attention.parameters(), reference.parameters(), strict=True): + assert actual_param.grad is not None and expected_param.grad is not None + torch.testing.assert_close(actual_param.grad.float(), expected_param.grad.float(), rtol=1e-5, atol=1e-6) + + def test_dspark_attention_casts_rotary_values_to_query_dtype(): attention = DSparkParallelAttention( hidden_size=256, diff --git a/vime/backends/megatron_utils/dspark/attention.py b/vime/backends/megatron_utils/dspark/attention.py index b14043b0..10c1c82f 100644 --- a/vime/backends/megatron_utils/dspark/attention.py +++ b/vime/backends/megatron_utils/dspark/attention.py @@ -47,7 +47,7 @@ def rotate_half(x): class DSparkRotaryEmbedding(nn.Module): - """Precompute rotary sin/cos for DSpark positions. + """Compute rotary sin/cos for DSpark positions. DSpark position ids cover both context (0..seq_len-1) and draft tokens (anchor_pos..anchor_pos+block_size-1 per block). The rotary table must @@ -58,11 +58,6 @@ def __init__(self, head_dim: int, rotary_base: float = 10000.0): super().__init__() self.head_dim = head_dim self.rotary_base = rotary_base - # Precompute a large table; will be indexed as needed. - # Max position ~ seq_len + num_anchors * block_size, which is bounded - # by the model's max_position_embeddings. - inv_freq = 1.0 / (rotary_base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) - self.register_buffer("inv_freq", inv_freq, persistent=False) def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Compute (cos, sin) for the given position ids. @@ -75,10 +70,16 @@ def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tenso # inv_freq: [head_dim/2] # position_ids: [bsz, seq_len] # freqs: [bsz, seq_len, head_dim/2] - inv_freq = self.inv_freq.float() # [head_dim/2] + # Module dtype conversions also round floating-point buffers. Recompute + # in FP32 on the input device rather than widening a rounded inv_freq. + inv_freq = 1.0 / ( + self.rotary_base + ** (torch.arange(0, self.head_dim, 2, device=position_ids.device, dtype=torch.float32) / self.head_dim) + ) positions = position_ids.float() # [bsz, seq_len] # Outer product per batch: [bsz, seq_len, head_dim/2] - freqs = torch.einsum("i,bj->bji", inv_freq, positions) + # Elementwise multiplication stays FP32 even under CUDA autocast. + freqs = positions.unsqueeze(-1) * inv_freq emb = torch.cat([freqs, freqs], dim=-1) # [bsz, seq_len, head_dim] cos = emb.cos() sin = emb.sin()