Conversation
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 <luozijian0924@gmail.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
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.
| cos = cos.unsqueeze(1) | ||
| sin = sin.unsqueeze(1) | ||
| return cos.to(position_ids.dtype), sin.to(position_ids.dtype) | ||
| return cos, sin |
There was a problem hiding this comment.
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))There was a problem hiding this comment.
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.
| torch.testing.assert_close(cos, expected_cos, rtol=0, atol=0) | ||
| torch.testing.assert_close(sin, expected_sin, rtol=0, atol=0) |
There was a problem hiding this comment.
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).
| 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) |
There was a problem hiding this comment.
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.
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 <luozijian0924@gmail.com>
|
Supplementary red/green check on unchanged head
This was a fresh assistant-executed Linux CPU run (Python 3.13.5, Torch 2.10.0+cpu, pytest 9.0.2). Both production files and the unchanged test file were checked against their Git blob hashes. The real attention implementation and FP16/BF16 forward/backward tests ran in an isolated two-file namespace tree, without full-package initialization, repository conftest, Megatron, or CUDA. It is supplementary source-level evidence, not an upstream/full-stack pass; the exact scope and hashes are in the description. Buildkite #1201 still reports |
1 participant
Summary
Fix DSpark rotary tables being cast to integer position-ID dtype, and preserve FP32 phases after parent-module dtype conversions and under autocast.
Original base:
ce92eff12ecdc81396bf41a2f94e62dd5b0aca32. Original PR commit:e9d427932. Current code HEAD:cfac5c2bacf6dcfe72c96b807a06c2da3d622dbe, appended without rewriting the original sign-off.Fresh supplementary CPU verification — September 16, 2026
The exact production attention module and committed test module were fetched from GitHub into an isolated two-file namespace tree. Their Git blob hashes were verified before execution:
attention.py10c1c82f3e861d35f3c857f4b9705ee308f51676test_dspark_rotary_embedding.pycbb552227df30140f440383e454b7768182cc188attention.pyatce92eff122378260b4593289b2221337fb144a43bc764e764The same unchanged test file was run against each production source in a separate process:
ce92eff12cfac5c2Environment: Linux, Python 3.13.5, PyTorch
2.10.0+cpu, pytest 9.0.2,OMP_NUM_THREADS=1,MKL_NUM_THREADS=1, andPYTEST_DISABLE_PLUGIN_AUTOLOAD=1.Scope: real PyTorch and the complete attention module, including the committed FP16/BF16 attention forward/backward tests. This isolated snapshot does not include the full VIME checkout, package initializers, repository conftest/configuration, Megatron installation, or GPU execution. No implementation or assertions were mocked or altered. This supplements, rather than replaces, full-repository CI. Source-hash checks, pytest output, and JUnit reports were retained locally; they are not uploaded CI artifacts.
Previously reported validation
The existing WSL run on the current source reported 14 passed, 2 CUDA-only skips with Python 3.12.3, PyTorch
2.11.0+cpu, and pytest 9.1.1. Applicable changed-file pre-commit hooks andgit diff --checkpassed at publication. These are historical results in a different environment, not an additional test total or newly rerun hooks.Retained development controls:
2.3.0a0+ebedce2, CUDA runtime 12.3 tested 24 combinations: FP16/BF16 modules, 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 earlier implementation matched the FP32 reference in 0/24 combinations; the final implementation matched in 24/24. Maximum observed candidate error was 0.0 in this finite probe, not a universal guarantee.The retained A100 environment predates
nn.RMSNorm, so that GPU probe loaded the exact RoPE class source, not full attention or an end-to-end training stack. It was not rerun for this update. The fresh CPU attention tests above must not be described as GPU attention validation.Upstream CI and review — checked September 16, 2026
Buildkite #1201 has GitHub status
success, but its detailed description ispassed and blocked. The repository separates always-on CPU jobs from a manually gated GPU stage; the green summary does not establish full GPU-suite execution.On this HEAD, the existing CPU pipeline calls
python -m pytest tests/test_dspark_rotary_embedding.py -q. CPU workers legitimately skip the two CUDA-only cases. Full GPU attention forward/backward remains an explicit validation gap, not a passing claim.The earlier automated review items are implemented in the current source. Maintainer review and confirmation of any required GPU execution remain outstanding. No gate or permission changes are proposed.
Scope
Three changed 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 dependency changes are introduced.
No model-quality improvement, distributed validation, full-training convergence, or whole-repository pass is claimed. This description/evidence update changes no source files, commit history, or existing sign-offs.
AI assistance: ChatGPT assisted with implementation, tests, verification, and PR drafting. The September 16 supplementary run was assistant-executed and is not represented as an independent human review.