[fsdp, recipe] feat: support BF16 LoRA training for Qwen3.8-Flash-Next

scripts/run_qwen38_lora_dapo.sh

@@ -0,0 +1,86 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Run in the prepared uv environment on an existing two-node Ray cluster.
+# Both engines use the same BF16 checkpoint derived from the FP8 source.
+# Model-specific Transformers/vLLM support must be installed separately.
+# Set VERL_NO_PLACEMENT_MMAP_DIR on Ray workers to node-local disk shared by that node's ranks.
+: "${MODEL_PATH:?set the shared BF16 checkpoint}"
+: "${TRAIN_FILE:?set the DAPO-Math training parquet}"
+: "${VAL_FILE:?set the held-out AIME2024 parquet}"
+: "${OUTPUT_DIR:?set the checkpoint directory}"
+
+exec uv run --active --no-sync python -m verl.trainer.main_ppo \
+    algorithm.adv_estimator=grpo \
+    algorithm.use_kl_in_reward=False \
+    data.train_files="['${TRAIN_FILE}']" \
+    data.val_files="['${VAL_FILE}']" \
+    data.train_batch_size=8 \
+    data.max_prompt_length=4096 \
+    data.max_response_length=4096 \
+    data.filter_overlong_prompts=True \
+    data.truncation=error \
+    actor_rollout_ref.model.path="$MODEL_PATH" \
+    actor_rollout_ref.model.trust_remote_code=True \
+    +actor_rollout_ref.model.override_config.attn_implementation=sdpa \
+    actor_rollout_ref.model.use_remove_padding=False \
+    actor_rollout_ref.model.enable_gradient_checkpointing=True \
+    actor_rollout_ref.model.use_liger=False \
+    actor_rollout_ref.model.use_fused_kernels=False \
+    actor_rollout_ref.model.lora_rank=16 \
+    actor_rollout_ref.model.lora_alpha=32 \
+    'actor_rollout_ref.model.target_modules=[q_proj,k_proj,v_proj,o_proj,in_proj_a,in_proj_b,in_proj_qkv,in_proj_z,out_proj,gate_proj,up_proj,down_proj]' \
+    '+actor_rollout_ref.model.target_parameters=[mlp.experts.gate_up_proj,mlp.experts.down_proj]' \
+    actor_rollout_ref.model.lora.merge=False \
+    actor_rollout_ref.model.mtp.enable=True \
+    actor_rollout_ref.model.mtp.enable_train=False \
+    actor_rollout_ref.model.mtp.enable_rollout=True \
+    actor_rollout_ref.model.mtp.method=mtp \
+    actor_rollout_ref.model.mtp.num_speculative_tokens=3 \
+    actor_rollout_ref.actor.strategy=fsdp2 \
+    +actor_rollout_ref.actor.checkpoint.save_lora_only=True \
+    actor_rollout_ref.actor.fsdp_config.model_dtype=bf16 \
+    actor_rollout_ref.actor.fsdp_config.fsdp_size=-1 \
+    actor_rollout_ref.actor.fsdp_config.param_offload=False \
+    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
+    actor_rollout_ref.actor.optim.lr=1e-5 \
+    'actor_rollout_ref.actor.optim.betas=[0.9,0.98]' \
+    actor_rollout_ref.actor.optim.weight_decay=0.1 \
+    actor_rollout_ref.actor.clip_ratio_high=0.28 \
+    actor_rollout_ref.actor.ppo_mini_batch_size=8 \
+    actor_rollout_ref.actor.use_dynamic_bsz=True \
+    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=8192 \
+    actor_rollout_ref.actor.use_kl_loss=False \
+    actor_rollout_ref.rollout.name=vllm \
+    actor_rollout_ref.rollout.dtype=bfloat16 \
+    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \
+    actor_rollout_ref.rollout.gpu_memory_utilization=0.76 \
+    actor_rollout_ref.rollout.load_format=safetensors \
+    actor_rollout_ref.rollout.n=8 \
+    actor_rollout_ref.rollout.seed=17 \
+    actor_rollout_ref.rollout.temperature=1.0 \
+    actor_rollout_ref.rollout.val_kwargs.do_sample=False \
+    actor_rollout_ref.rollout.val_kwargs.n=1 \
+    actor_rollout_ref.rollout.max_model_len=8192 \
+    actor_rollout_ref.rollout.max_num_batched_tokens=8192 \
+    actor_rollout_ref.rollout.max_num_seqs=64 \
+    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \
+    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=8192 \
+    actor_rollout_ref.rollout.enforce_eager=False \
+    "+actor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config='{\"inductor_compile_config\":{\"triton.autotune_at_compile_time\":false}}'" \
+    +actor_rollout_ref.rollout.engine_kwargs.vllm.language_model_only=True \
+    +actor_rollout_ref.rollout.engine_kwargs.vllm.enable_flashinfer_autotune=False \
+    +actor_rollout_ref.rollout.engine_kwargs.vllm.gdn_prefill_backend=triton \
+    actor_rollout_ref.rollout.checkpoint_engine.backend=naive \
+    trainer.logger='["console"]' \
+    trainer.project_name=qwen38_lora \
+    trainer.experiment_name=bf16_fsdp16_tp4 \
+    trainer.n_gpus_per_node=8 \
+    trainer.nnodes=2 \
+    trainer.total_training_steps=200 \
+    trainer.val_before_train=True \
+    trainer.save_freq=10 \
+    trainer.test_freq=10 \
+    trainer.default_local_dir="$OUTPUT_DIR" \
+    +ray_kwargs.ray_init.address="${RAY_ADDRESS:-auto}" \
+    "$@"

tests/special_distributed/run_all.sh

@@ -25,3 +25,4 @@ torchrun --nproc-per-node=2 --standalone tests/special_distributed/test_fsdp2_cp
 # Regression for colocated FSDP2 full-parameter model transfers. It locks in
 # PyTorch's implicit pinned allocation for non-blocking D2H Module.to().
 torchrun --nproc-per-node=2 --standalone tests/special_distributed/test_fsdp2_pinned_model_transfer.py
+torchrun --nproc-per-node=2 --standalone tests/special_distributed/test_fsdp2_full_state_load.py

tests/special_distributed/test_fsdp2_full_state_load.py

@@ -0,0 +1,205 @@
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Regression test for FSDP2 loading from a rank-zero CPU full state.
+
+Launch:
+    torchrun --standalone --nproc-per-node=2 \
+        tests/special_distributed/test_fsdp2_full_state_load.py
+"""
+
+import os
+import tempfile
+from pathlib import Path
+from types import SimpleNamespace
+
+import torch
+import torch.distributed as dist
+from torch import nn
+from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict
+from torch.distributed.device_mesh import init_device_mesh
+from transformers import Qwen2Config, Qwen2ForCausalLM
+
+from verl.utils.distributed import initialize_global_process_group
+from verl.utils.fsdp_utils import (
+    MixedPrecisionPolicy,
+    apply_fsdp2,
+    fsdp2_load_full_state_dict,
+    get_no_placement_param_registrations,
+    materialize_no_placement_params,
+    set_no_placement_param_registrations,
+    temporarily_detach_no_placement_params,
+    to_empty_preserving_shared_params,
+)
+
+
+class ToyBlock(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.embedding = nn.Embedding(8, 4)
+        self.linear = nn.Linear(4, 4, bias=False)
+
+    def forward(self, tokens):
+        return self.linear(self.embedding(tokens))
+
+
+class BufferedModel(nn.Module):
+    _no_split_modules = ["ToyBlock"]
+    _no_placement_params = ["large.weight"]
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.config = SimpleNamespace(tie_word_embeddings=False)
+        self.block = ToyBlock()
+        self.large = nn.Embedding(8, 4)
+        self.large.weight.requires_grad_(False)
+        self.register_buffer("marker", torch.arange(4, dtype=torch.bfloat16), persistent=False)
+
+    def forward(self, tokens):
+        return self.block(tokens)
+
+
+def _build_model(rank: int) -> tuple[BufferedModel, dict[str, torch.Tensor]]:
+    model = BufferedModel()
+    if rank == 0:
+        model.large.weight.copy_(torch.arange(32, dtype=torch.float32).view(8, 4))
+    else:
+        model.large.weight = nn.Parameter(
+            torch.empty(model.large.weight.shape, device="meta"),
+            requires_grad=False,
+        )
+    registrations = materialize_no_placement_params(get_no_placement_param_registrations(model))
+    set_no_placement_param_registrations(model, registrations)
+    with temporarily_detach_no_placement_params(model):
+        state = model.state_dict() if rank == 0 else {}
+    return model, state
+
+
+def _forbid_full_model_to(*args, **kwargs) -> None:
+    raise AssertionError("full model must not be materialized with Module.to")
+
+
+def _check_tied_checkpoint_roundtrip(rank, mesh):
+    config = Qwen2Config(
+        vocab_size=32,
+        hidden_size=16,
+        intermediate_size=32,
+        num_hidden_layers=1,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        tie_word_embeddings=True,
+    )
+    model = Qwen2ForCausalLM(config).to(dtype=torch.bfloat16)
+    state = model.state_dict() if rank == 0 else {}
+    buffers = {name: buffer.detach().cpu() for name, buffer in model.named_buffers()}
+    to_empty_preserving_shared_params(model, "meta")
+    assert model.model.embed_tokens.weight is model.lm_head.weight
+    apply_fsdp2(
+        model,
+        {"mesh": mesh, "mp_policy": MixedPrecisionPolicy(param_dtype=torch.bfloat16), "reshard_after_forward": True},
+        {"wrap_policy": {"transformer_layer_cls_to_wrap": ["Qwen2DecoderLayer"]}},
+    )
+    fsdp2_load_full_state_dict(model, state, mesh, buffers=buffers)
+    assert model.model.embed_tokens.weight is model.lm_head.weight
+    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
+    tokens = torch.tensor([[1, 2, 3, 4]], device="cuda")
+    model(input_ids=tokens, labels=tokens).loss.backward()
+    optimizer.step()
+    merged = get_model_state_dict(model, options=StateDictOptions(full_state_dict=True, cpu_offload=True))
+    if rank == 0:
+        torch.testing.assert_close(merged["model.embed_tokens.weight"], merged["lm_head.weight"], atol=0, rtol=0)
+        with tempfile.TemporaryDirectory() as directory:
+            Qwen2ForCausalLM(config).to(dtype=torch.bfloat16).save_pretrained(directory, state_dict=merged)
+            restored = Qwen2ForCausalLM.from_pretrained(directory, dtype=torch.bfloat16).state_dict()
+            for name in merged:
+                torch.testing.assert_close(merged[name], restored[name], atol=0, rtol=0)
+    dist.barrier()
+
+
+def main() -> None:
+    if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
+        print("test_fsdp2_full_state_load skipped: requires two CUDA devices")
+        return
+    _, rank, world_size = initialize_global_process_group()
+    if world_size != 2:
+        raise RuntimeError(f"expected two ranks, got {world_size}")
+    torch.cuda.set_device(rank)
+    mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("fsdp",))
+    with tempfile.TemporaryDirectory() as directory:
+        directories = [directory]
+        dist.broadcast_object_list(directories, src=0)
+        os.environ["VERL_NO_PLACEMENT_MMAP_DIR"] = directories[0]
+        model, full_state = _build_model(rank)
+        ptr = model.large.weight.data_ptr()
+        mapping = next(
+            line.split()
+            for line in Path("/proc/self/maps").read_text().splitlines()
+            if int(line.split()[0].split("-")[0], 16) <= ptr < int(line.split()[0].split("-")[1], 16)
+        )
+        backing_files = [None] * world_size
+        dist.all_gather_object(backing_files, mapping[3:5])  # device and inode, not virtual address
+        assert backing_files[0] == backing_files[1]
+        assert mapping[4] != "0"
+        dist.barrier()
+        assert not list(Path(directories[0]).iterdir())
+        dist.barrier()
+        del os.environ["VERL_NO_PLACEMENT_MMAP_DIR"]
+    expected = (
+        full_state["block.linear.weight"].detach().clone().to("cuda")
+        if rank == 0
+        else torch.empty((4, 4), device="cuda", dtype=torch.float32)
+    )
+    dist.broadcast(expected, src=0)
+    buffers = {name: buffer.detach().cpu() for name, buffer in model.named_buffers()}
+    registrations = vars(model)["_verl_no_placement_param_registrations"]
+    with temporarily_detach_no_placement_params(model):
+        to_empty_preserving_shared_params(model, "meta")
+    apply_fsdp2(
+        model,
+        {
+            "mesh": mesh,
+            "mp_policy": MixedPrecisionPolicy(
+                param_dtype=torch.bfloat16,
+                reduce_dtype=torch.float32,
+                cast_forward_inputs=True,
+            ),
+            "offload_policy": None,
+            "reshard_after_forward": True,
+            "ignored_params": {param for _, _, param, _ in registrations},
+        },
+        {"wrap_policy": {"transformer_layer_cls_to_wrap": ["ToyBlock"]}},
+    )
+    model.to = _forbid_full_model_to
+    with temporarily_detach_no_placement_params(model):
+        fsdp2_load_full_state_dict(model, full_state, mesh, buffers=buffers)
+
+    torch.testing.assert_close(model.block.linear.weight.full_tensor().float(), expected, atol=0.0, rtol=0.0)
+    torch.testing.assert_close(model.marker, torch.arange(4, device="cuda", dtype=torch.bfloat16))
+    assert model.large.weight.device.type == "cpu"
+    expected_large = torch.arange(32, dtype=torch.float32).view(8, 4)
+    torch.testing.assert_close(model.large.weight, expected_large)
+    # A nested embedding must not create an overlapping FSDP unit inside ToyBlock.
+    output = model(torch.tensor([1, 2], device="cuda"))
+    output.sum().backward()
+    assert torch.isfinite(output).all()
+    assert model.block.linear.weight.grad is not None
+    _check_tied_checkpoint_roundtrip(rank, mesh)
+    dist.barrier()
+    dist.destroy_process_group()
+    if rank == 0:
+        print("test_fsdp2_full_state_load passed")
+
+
+if __name__ == "__main__":
+    main()

tests/test_protocol_v2_on_cpu.py

@@ -149,6 +149,36 @@ def test_index_select_tensor_dict_preserves_3d_nested_tensor_layout_with_equal_s
     tu.assert_tensordict_eq(selected, tu.get_tensordict({"position_ids": expected}))
 
 
+def test_fix_3d_position_ids_rebuilds_equal_length_layout():
+    rows = [torch.arange(20).reshape(4, 5), torch.arange(20, 40).reshape(4, 5)]
+    position_ids = torch.nested.as_nested_tensor(rows, layout=torch.jagged)
+    data = tu.get_tensordict({"position_ids": position_ids})
+
+    tu.maybe_fix_3d_position_ids(data)
+
+    assert data["position_ids"]._ragged_idx == 2
+    assert data["position_ids"].values().shape == torch.Size([4, 10])
+    assert data["position_ids"].offsets().tolist() == [0, 5, 10]
+    assert torch.equal(data["position_ids"][0], rows[0])
+    assert torch.equal(data["position_ids"][1], rows[1])
+
+
+def test_fix_3d_position_ids_restores_lost_ragged_index():
+    position_ids = tu.nested_tensor_from_tensor_list(
+        [torch.arange(8).reshape(4, 2), torch.arange(20).reshape(4, 5)], ragged_idx=2
+    )
+    values_ptr = position_ids.values().data_ptr()
+    offsets_ptr = position_ids.offsets().data_ptr()
+    position_ids._ragged_idx = 1
+    data = tu.get_tensordict({"position_ids": position_ids})
+
+    tu.maybe_fix_3d_position_ids(data)
+
+    assert data["position_ids"]._ragged_idx == 2
+    assert data["position_ids"].values().data_ptr() == values_ptr
+    assert data["position_ids"].offsets().data_ptr() == offsets_ptr
+
+
 def test_tensordict_with_images():
     # each sample contains a sequence with multiple images of different sizes
     vocab_size = 128