[Bugfix][Training] Preserve DSpark rotary embedding dtype

[LOGO127]

Summary

Fix DSpark rotary tables being cast to integer position-ID dtype, and preserve FP32 phases after parent-module dtype conversions and under autocast.

• Keep cosine/sine floating point; convert them to the query dtype only at the attention boundary.
• Recompute inverse frequencies directly in FP32 on the positions device instead of widening an already-rounded FP16/BF16 buffer.
• Use an elementwise outer product to avoid CUDA autocast downcasting the phase computation.
• Cover dtype casts and round trips, long positions, CPU/CUDA autocast, and attention forward/backward. Use explicit small nonzero tolerances.
• Invoke the test file through pytest in the CPU pipeline. The previous direct Python invocation imported the definitions without running these tests.

Original base: ce92eff12ecdc81396bf41a2f94e62dd5b0aca32. Original PR commit: e9d427932. Signed follow-up: cfac5c2bacf6dcfe72c96b807a06c2da3d622dbe, appended without rewriting the original signed commit.

Validation

On the exact signed follow-up, freshly rerun:

python -m pytest tests/test_dspark_rotary_embedding.py -q

14 passed, 2 skipped on WSL Linux, Python 3.12.3, PyTorch 2.11.0+cpu, pytest 9.1.1. The two skips explicitly require CUDA. CPU tests include actual FP16/BF16 attention outputs, input gradients and parameter gradients. All applicable changed-file pre-commit hooks and git diff --check passed; the signed checkout remains unchanged.

Retained red/green evidence from developing the same file tree:

• The original main returned int64 rotary tables for long position IDs, including an all-zero sine table on the A100 probe.
• The initial 12-case review regression suite gave 10 failures / 2 passes on the previously published implementation.
• A source-level RoPE probe on NVIDIA A100 80GB PCIe, PyTorch 2.3.0a0+ebedce2, CUDA runtime 12.3 tested 24 combinations: module FP16/BF16, with/without conversion back to FP32, head-dimension/base pairs (64, 10000) and (128, 500000), and autocast off/FP16/BF16, including positions up to 32769. The previous implementation matched the FP32 reference in 0/24 combinations; the final implementation matched in 24/24. The observed maximum candidate error in this finite probe was 0.0, not a guarantee for every configuration.
• Recomputing frequencies alone was insufficient in the intermediate A100 probe; the elementwise product addresses the remaining autocast difference.

The signed follow-up has exactly the validated file tree. A100 evidence was retained, not rerun during publication. The rented host's PyTorch predates nn.RMSNorm, so that GPU probe loads the exact RoPE class source, not full attention or an end-to-end training stack. No training-quality improvement, distributed validation or whole-repository pass is claimed.

Scope and review

Three files: rotary implementation, focused regression module, and its existing CPU premerge route. The removed inverse-frequency buffer was non-persistent, so no persistent checkpoint keys are removed. No rollout protocol, optimizer behavior or new dependency is introduced. Both automated review threads are addressed in the follow-up; maintainer review remains pending.

AI assistance: ChatGPT assisted with implementation, tests, validation, and PR drafting.

[gemini-code-assist]

Code Review

This pull request modifies the DSparkRotaryEmbedding to return cos and sin tensors in their original precision rather than casting them to the integer-based position_ids.dtype. It also updates DSparkParallelAttention to cast these embeddings to the query's data type (q.dtype) and adds corresponding unit tests to the Buildkite pipeline. The review feedback highlights two key improvement opportunities: first, addressing a potential precision loss in inv_freq when the parent module is cast to a lower precision (such as bfloat16) by computing it dynamically in float32; second, avoiding strict zero tolerances (rtol=0, atol=0) in floating-point assertions to prevent flaky tests across different hardware architectures.

[gemini-code-assist]

While returning cos and sin in float32 prevents immediate casting to position_ids.dtype, there is a subtle precision issue with the inv_freq buffer.

When the parent module is cast to a lower precision (e.g., via attention.to(dtype=torch.bfloat16)), PyTorch recursively casts all registered buffers, including self.inv_freq, to bfloat16. In forward(), calling self.inv_freq.float() upcasts it back to float32, but the precision lost during the downcast to bfloat16 is already gone. This can lead to significant phase errors when multiplied by large position IDs in long-context scenarios.

To prevent this, you can compute inv_freq dynamically in forward() directly in float32 on the correct device, which also avoids any device mismatch issues:

inv_freq = 1.0 / (self.rotary_base ** (torch.arange(0, self.head_dim, 2, device=position_ids.device, dtype=torch.float32) / self.head_dim))

[LOGO127]

Addressed in signed follow-up cfac5c2b. Inverse frequencies are now generated directly in FP32 on the positions device, rather than recovered from a dtype-cast buffer. The phase outer product uses elementwise multiplication: the first attempt that only regenerated frequencies still diverged under CUDA autocast.

Added parent FP16/BF16 casts and cast-back, long-position, and attention forward/backward regressions. On this signed tree the CPU suite reran with 14 passed / 2 CUDA-only skips. Retained A100 RoPE-only checks on the same implementation passed 24/24 combinations, versus 0/24 for the prior implementation; full GPU attention/training was not run. The CPU CI now invokes pytest instead of only importing test definitions. Commands and environment limits are in the updated PR body.

AI assistance: ChatGPT assisted implementation, validation, and this response.

[gemini-code-assist]

Using strict zero tolerances (rtol=0, atol=0) for floating-point assertions can lead to flaky tests across different hardware architectures, compiler optimizations, or PyTorch versions. It is safer to use PyTorch's default tolerances for float32 by omitting these arguments, or specify a small non-zero tolerance (e.g., atol=1e-7, rtol=1e-7).

Suggested change

	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)
	torch.testing.assert_close(sin, expected_sin)

[LOGO127]

Addressed in cfac5c2b: floating-point comparisons now use explicit rtol=1e-5, atol=1e-6 rather than zero tolerances. The integer-output and rounded-frequency bugs remain covered by the regression controls. The exact signed CPU rerun collected 16 cases: 14 passed and 2 explicitly CUDA-only skips. No change to production acceptance thresholds or training tolerances.

AI assistance: ChatGPT assisted with the tests and this response.

[read-the-docs-community]

Documentation build overview

📚 vime | 🛠️ Build #34501744 | 📁 Comparing cfac5c2 against latest (ce92eff)

  🔍 Preview build  

45 files changed · ± 45 modified

± Modified

• genindex.html
• index.html
• search.html
• advanced/arch-support-beyond-megatron.html
• advanced/delta-weight-sync.html
• advanced/external-rollout-engines.html
• advanced/fault-tolerance.html
• advanced/low-precision.html
• advanced/megatron-config.html
• advanced/observability.html
• and 35 more...