Fix rotary frequency precision

[le1nux]

Summary

RotaryTransform._update_cos_sin_tables cast the inverse frequencies to the activation dtype
before scaling them by token positions (self.inv_freq.to(x.dtype)), so under BF16 training the
rotary phase t * inv_freq was accumulated from BF16 frequencies. Because the phase grows with the
absolute token position, a 2^-9 relative error on a frequency becomes an O(1) error in radians at
long context.

• keep the inverse frequencies in FP32 while scaling by token positions
• query/key outputs are unchanged: only the phase and the trigonometric tables are computed in
  FP32, and cos/sin are still cast back to the activation dtype, exactly as
  LlamaRotaryEmbedding in src/modalities/conversion/gpt2/modeling_gpt2.py already did — that
  path forced FP32 all along, so the codebase disagreed with itself
• add a regression test that records the forward-pass operand and result dtypes
• align the YaRN test reference on FP32 inverse frequencies (no-op for that test, which runs in
  FP32, but it makes the reference independent of the input dtype)

Impact

Max phase error of the cos/sin tables against an FP64 reference (n_embd=4096, n_head=32,
base_freq=10000, BF16 activations):

seq_len	before	after
512	6.2e-01 rad	2.8e-03 rad
2048	2.5e+00 rad	2.8e-03 rad
4096	3.14e+00 rad (π)	2.8e-03 rad
8192	3.14e+00 rad (π)	2.9e-03 rad

At 4k tokens and beyond the positional phase was effectively meaningless. Every shipped FSDP2
config reaches this code via qkv_transforms: RotaryTransform.

The remaining 2.8e-03 rad is the irreducible cost of storing cos/sin in BF16, and it is the
same residual the reference implementations carry: after this change the tables are bit-identical
to LlamaRotaryEmbedding with FP32 inverse frequencies (max |diff| = 0.0e+00 for both cos and
sin at seq 4096). The training path now computes standard RoPE, so HF export and vLLM inference
compute the same positions as training.

Effect on existing FSDP2 checkpoints

The fix applies in full when warm-starting an older checkpoint. FSDP2's MixedPrecisionPolicy has
no buffer_dtype and does not cast buffers, so inv_freq stayed FP32 throughout every FSDP2 run
and the stored value was never damaged — the bug was only in the einsum. Checked against a real
pre-fix DCP checkpoint (32 inv_freq keys, all torch.float32): the loaded values match freshly
computed FP32 frequencies to 3.0e-08, i.e. ~1 FP32 ulp, and are not BF16-rounded (which would
have cost 1.8e-03). Warm-starting that checkpoint gives 3.1e-03 rad at seq 8192, against π before.
Loading is in-place into the model's own FP32 buffer via StateDictOptions(strict=False)
(app_state.py:206), so a checkpoint without inv_freq keys is equally fine — the buffer keeps
its freshly computed FP32 value.

Worth being deliberate about: this changes the positional encoding a resumed run sees. A run that
trained under the old, self-consistent-but-nonstandard phase and then resumes with standard RoPE is
undergoing a mid-training change to the position function, so expect a loss bump on resume. An
anneal or an export is the natural moment to take the correction.

Known limitation: FSDP1

.float() can only preserve precision while the buffer is still FP32. Two paths cast inv_freq
to BF16 before it reaches this line, and upcasting cannot recover the discarded mantissa bits, so
the change is a no-op there:

• FSDP1 with mixed_precision_settings: BF_16 — bfSixteen sets buffer_dtype=torch.bfloat16
  (env_utils.py:43-49) and FSDP1 casts buffers on every root pre-forward. Measured: still
  3.14e+00 rad at seq 4096. Reaches 5 shipped configs that pair fsdp1_wrapped + BF_16 +
  RotaryTransform. MIXED_PRECISION_MEGATRON leaves buffer_dtype unset and is unaffected.
• fsdp1_checkpointed with precision: BF16 — TorchCheckpointLoading.load_model_checkpoint
  does model.to(self.device, dtype=...) (torch_checkpoint_loading.py:44), which casts the
  buffer at load time regardless of what the file holds. Measured: 1.2e-01 rad at only seq 128 with
  base_freq=500000. Same root cause in conversion_model.py:25, which builds the HF model with
  .to(dtype=torch.bfloat16); note that check_converted_model then compares two equally-affected
  models, and that exported checkpoints themselves are fine because inv_freq is persistent=False
  on the HF side, so transformers and vLLM recompute it in FP32.

This is left as-is deliberately: FSDP1 has carried a FutureWarning since 0.4
(model_factory.py:112) and is not planned to be supported going forward. Closing it properly
means not trusting the buffer's dtype at all — regenerating the frequencies when
self.inv_freq.dtype != torch.float32 rather than upcasting (verified to bring both paths to
2.8e-03 rad) — which is more surgery than a deprecated path warrants.

Related: inv_freq is a persistent buffer, so a checkpoint written by an FSDP1 + BF_16 run
stores BF16-rounded frequencies and load_state_dict casts rather than rejects them
(max |loaded - correct| = 1.2e-03, exactly BF16-rounded). Such a checkpoint stays affected after
this fix. FSDP2 + dcp checkpoints are unaffected, as above.

Validation

• pytest tests (excluding end2end_tests, single node): 497 passed, 44 skipped, no regressions
• pytest tests/conversion tests/test_rotary_qkv_transform.py tests/nn: 37 passed
• the new regression test fails on the pre-fix source and passes on the fixed source
• pre-commit gate at the pinned versions (isort 5.11.5, black 23.9.1, ruff 0.0.278): clean