fix(megatron): isolate FP32 adapter optimizer state

[bvolpato]

Summary

With Megatron multi-LoRA and trainer.bf16=false, switching adapters restores the model weights but leaves the previous adapter's Adam moments live. A fresh adapter can therefore update in the opposite direction to its own gradient.

Megatron keeps native FP32 optimizer parameters in shard_fp32_groups; AdapterStore currently snapshots only shard_fp32_from_float16_groups. Use one iterator over both layouts for slot allocation, snapshot, and restore. Existing BF16 masters and chained optimizers retain their ordering.

Testing

Eight CPU numerical regressions compare interleaved adapters against separately trained Adam baselines: FP32, BF16, mixed groups, and chained optimizers, each with per-parameter or TE-style group counters. The tests use the actual AdapterStore and real PyTorch Adam updates; CUDA container types and pinning are replaced for CPU execution.

All eight tests pass. Against the unchanged production module, the six FP32/mixed/chained cases fail while both BF16 controls pass. Executed with PyTorch 2.11.0+cpu through isolated uv; an independent review also ran all eight successfully. Required pre-commit hooks pass. The existing lifecycle module skips without Megatron; GPU/Transformer Engine execution remains unverified.

The GitHub GPU job exited before running tests: ANYSCALE_CLI_TOKEN was empty and Anyscale rejected the credentials.

Repro command

From the repository root:

uv run --isolated --extra skyrl-train --extra dev pytest --noconftest \
  tests/backends/skyrl_train/workers/test_adapter_store_optimizer_state.py -q

Downsides

Each adapter slot, including the pristine template, now stores native FP32 parameter shards and their optimizer state. For ordinary FP32 Adam this adds about 12 bytes per local FP32 parameter element, plus scalar state, and corresponding copies during swaps. The parameter copy duplicates values already in the DDP buffer snapshot; keeping the existing slot layout avoids a separate state-only path. BF16 parameter storage is unchanged.

Risk and rollback

This affects Megatron multi-LoRA swaps. Incorrect group ordering would show up as adapter-specific parameter or moment mismatches, which the mixed and chained numerical tests cover. Revert the commit and restart workers from saved checkpoints to roll back; the saved checkpoint format is unchanged.