Fix meta-init bf16 cast corrupting rotary inv_freq under sequence parallelism

[jamesbraza]

Meta-init cast non-persistent buffers (rotary inv_freq) to bf16 while rank-0's from_pretrained kept them fp32, so the init-time buffer broadcast reinterpreted rank-0's fp32 bytes as bf16 garbage and produced NaN attention under SP>1. Now we cast only params and persistent buffers, matching transformers from_pretrained.

Closes #1709

[gemini-code-assist]

Code Review

This pull request modifies the meta-initialization logic in model_wrapper.py to cast only parameters and persistent buffers to the target dtype, preventing corruption of non-persistent buffers (like inv_freq in rotary embeddings) during FSDP initialization. It also adds corresponding unit and integration tests. The reviewer identified critical issues, including the use of a non-existent PyTorch method named_non_persistent_buffers() which will cause runtime AttributeErrors, the use of a non-existent model (Qwen/Qwen3-0.6B) in tests, and a similar blanket .to cast issue in fsdp_worker.py that needs to be addressed.

[SumanthRH]

@jamesbraza thanks for the PR!

As a sanity check I ran the test script test_meta_init.py:

uv run --isolated --extra dev --extra fsdp pytest -s -vvv tests/backends/skyrl_train/gpu/gpu_ci/test_meta_init.py

against origin/main

I expected NaN results without the fix in your PR, but tests pass. Can you confirm that the unit test is correctly simulating the issue you saw?

[jamesbraza]

I expected NaN results without the fix in your PR, but tests pass. Can you confirm that the unit test is correctly simulating the issue you saw?

Yep sorry about that, you're right the GPU tests were passing on main. Fixed it just now, turns out there was a small detail missing. Nice catch

[SumanthRH]

@jamesbraza thanks for the updated test, but it would be great to get a regression test for the NaN failures with SP > 1.

I re-ran the updated test and

1. assertions fail for dtype - inv_freq gets cast to bfloat16 instead of float32
2. assertions pass for Nan - no NaNs observed with the model used.

The reproduction script you provided before in #1709 clearly produces NaNs with BF16, but then the script forces meta init for Qwen3-1.7B, which clearly uses tie_word_embeddings=True, which means that SkyRL code skips meta init: https://huggingface.co/Qwen/Qwen3-1.7B/blob/main/config.json

So the reproducer is not simulating an actual FSDP init for Qwen3 1.7B with SkyRL.

[jamesbraza]

@jamesbraza thanks for the updated test, but it would be great to get a regression test for the NaN failures with SP > 1.
...

Yeah thanks for pointing this out, agreed let's get this right, I just pushed a commit upgrading the test.

So inv_freq holds head_dim/2 rotary frequencies. The prior model tiny-random-Llama-3 has head_dim of 4, so it only had 2 rotary frequencies. Now with Qwen3-8B (head_dim of 128) there's 64 rotary frequencies, so the NaN will be reliably reproduced (it worked on 3/3 attempts) on a non-rank-0 inv_freq.

Also Qwen3-8B is a non-tied model, so init_model now takes the meta path on its own, covering that code.

[SumanthRH]

Great catch and thanks for the hard work on the regression test.