refactor(meta-init): unify the meta-init build/restore path across bundles (fixes garbage RoPE on VeOmni / sharded-FSDP)

[zzhuoxin1508]

Summary

This started as a bug fix and turned into a refactor of the meta-init machinery.

Root cause (unchanged): under the meta-init path (meta_init_transformer=true,
used by the VeOmni / sharded-FSDP load_sharded backends) several bundles built the
transformer on torch.device("meta"), which forces non-persistent buffers onto
meta too — exactly the init-computed tables the checkpoint does not carry
(WanRotaryPosEmbed.freqs_cos/freqs_sin, SD3's sincos pos_embed, rope tables).
After to_empty() these hold garbage, and the no-op init_weights never rebuilds
them.

For WAN 2.1 this silently breaks RoPE: self-attn q/k get multiplied by junk cos/sin.
Because the multiply is by ~zero garbage, the gradient to the attn1.to_q /
attn1.to_k LoRA adapters is killed — those never train while attn2 / v /
out train normally, so LoRA-RL reward plateaus below the classic-FSDP baseline.

What changed (refactor)

The base already restored captured state via bundle._meta_init_state in
load_trainable_weights — but only bundles that actually captured benefited, and
the capture path was split across two mechanisms. This PR unifies it:

• unirl/models/types/meta_init.py
  • Add build_meta_init_transformer(factory, *, dtype): builds under
    accelerate.init_empty_weights(include_buffers=False) (params on meta,
    non-persistent buffers/attrs real on CPU), captures them via
    capture_init_state, casts dtype, and no-ops init_weights. Returns
    (transformer, captured).
  • Remove finalize_meta_init (folded into the above) and
    stamp_init_state_restore (the model-bound deferred-closure restore, which can
    be dropped across Ray-actor boundaries — the bundle carrier already covers it,
    robustly).
  • Keep capture_init_state / restore_init_state.
• All 7 meta-init bundles (wan21, sd3, qwen3, qwen_vl, flux2_klein,
  hunyuan_video, hunyuan_video15) now build via build_meta_init_transformer
  and stash the capture on bundle._meta_init_state.
• Restore is unchanged in placement: load_trainable_weights replays
  restore_init_state(model, bundle._meta_init_state) after the sharded weight load
  (Ray-robust — a plain picklable dict survives crossing actors).
• unirl/train/backend/sharded_load.py — comment cleanup only.

Net: 9 files, +128 / −190 — a single build/restore path, less code.

Related Issue

N/A

Test Plan

WAN 2.1-T2V-1.3B FlowGRPO, VeOmni backend, 720p, 2×8 GPU:
curve in comment
SD3.5-medium FlowGRPO, VeOmni backend, 8 GPU (diffusion/sd3_trainside_veomni):

• Ran end-to-end on the fix branch and confirmed reward grows normally

Compatibility / Risk

Behavioral change only on the meta_init_transformer=true path (VeOmni / sharded
FSDP). The classic-FSDP from_pretrained (eager) path is untouched. No
config/checkpoint/API format changes. Models with no init-computed non-persistent
buffers are unaffected — restore_init_state is a no-op when the capture is empty.

Reviewer Notes

AI-assisted. Root-caused and validated end-to-end on WAN 2.1 (reward alignment +
gradient map). The capturing build path lives in the shared
build_meta_init_transformer, so every meta-init model (e.g. SD3's sincos
pos_embed) gets the same restore for free.

Checklist

• I reviewed the changed code and removed unrelated/generated artifacts.
• I updated tests, docs, and configs where needed, or explained why not.

[zzhuoxin1508]