diff --git a/examples/train/fp8/README.md b/examples/train/fp8/README.md index 6ee2320439..4a2db41d48 100644 --- a/examples/train/fp8/README.md +++ b/examples/train/fp8/README.md @@ -30,8 +30,7 @@ Notes: `trainer.placement.policy_num_gpus_per_node` and the inference engines for a disaggregated placement. - **Qwen3.5 runs text-only.** All scripts set `language_model_only=true` on the policy, ref and - inference engine: Qwen3.5 otherwise loads through the VL bridge, which packs sequences inside - its own forward and is rejected together with SkyRL sample packing. + inference engine, so the trainer builds only the language model (no vision tower). - **GDN kernels on Blackwell.** The Blackwell scripts `export FLA_TILELANG=0` so fla uses its Triton GatedDeltaNet kernels; the TileLang packed backward aborts on B200 (it shows up as a CUDA "misaligned address" in the first backward). Leave it unset on Hopper, where the Triton diff --git a/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py b/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py index 391134ef67..39e08d812b 100644 --- a/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py +++ b/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py @@ -751,28 +751,30 @@ def preprocess_packed_seqs( return input_ids, packed_seq_params -def model_packs_sequences_internally(model: Union[nn.Module, List[nn.Module]]) -> bool: - """Whether the model packs sequences inside its own ``forward``. +def model_owns_vlm_packing(model: Union[nn.Module, List[nn.Module]]) -> bool: + """Whether the VLM handles a packed [1, T] stream itself. - True for ``Qwen3VLModel`` (e.g. Qwen3.5 via the VL bridge), which would - double-pack and corrupt the GDN ``cu_seqlens`` under SkyRL sample packing, so - :class:`MegatronModelWrapper` refuses packing for it. Returns ``False`` when - mbridge / Qwen3VL is not importable, so other models are unaffected. + True when every model chunk is Megatron-Bridge's ``Qwen3VLModel`` (Qwen3-VL, + Qwen3.5-VL), which rebuilds 3D mRoPE positions per packed sub-sequence from + ``packed_seq_params``, or sets ``model_owns_packing = True`` (NeMo-RL's opt-in + attribute for models that pack and split for context parallelism themselves). """ try: from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model import ( Qwen3VLModel, ) except ImportError: - return False + Qwen3VLModel = None chunks = model if isinstance(model, (list, tuple)) else [model] for chunk in chunks: unwrapped = unwrap_model(chunk) - unwrapped_list = unwrapped if isinstance(unwrapped, (list, tuple)) else [unwrapped] - if any(isinstance(m, Qwen3VLModel) for m in unwrapped_list): - return True - return False + if getattr(unwrapped, "model_owns_packing", False): + continue + if Qwen3VLModel is not None and isinstance(unwrapped, Qwen3VLModel): + continue + return False + return bool(chunks) def remove_left_padding( diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py index bc798bee64..02779eb32b 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py @@ -10,6 +10,7 @@ finalize_model_grads_with_expert_adapter_sync, ) from megatron.core.pipeline_parallel import get_forward_backward_func +from megatron.core.utils import unwrap_model from omegaconf import OmegaConf from skyrl.backends.skyrl_train.distributed.megatron.fused_lm_head import ( @@ -19,7 +20,7 @@ from skyrl.backends.skyrl_train.distributed.megatron.megatron_utils import ( get_model_config, make_batch_generator, - model_packs_sequences_internally, + model_owns_vlm_packing, preprocess_packed_seqs, recover_left_padding, remove_left_padding, @@ -215,24 +216,14 @@ def __init__( self.policy_loss_fn = policy_loss_fn self.remove_microbatch_padding = self.cfg.remove_microbatch_padding self.is_vlm = is_vlm + # Sample packing for VLMs needs a model that computes mRoPE positions from the packed + # stream itself; other VLMs only run unpacked (see _assert_vlm_supported). + self.model_owns_vlm_packing = is_vlm and model_owns_vlm_packing(self.actor_module) # Fuse the LM-head projection into the chunked log-prob/entropy via the # GPTModel output_processor hook (avoids materializing the full # [B, S, vocab//TP] logits + its fp32 grad). See model_utils. self._fused_lm_head = bool(getattr(self.cfg, "fused_lm_head_logprob", False)) self._fused_lm_head_backend = getattr(self.cfg, "fused_lm_head_logprob_backend", "torch") - # Some models (e.g. Qwen3.5 via the VL bridge -> Qwen3VLModel) pack - # sequences inside their own forward; SkyRL sample packing would then - # double-pack and corrupt the GDN cu_seqlens, so refuse it. For Qwen3.5, - # use language_model_only=True (native GPTModel GDN path) to pack. - if self.remove_microbatch_padding and model_packs_sequences_internally(self.actor_module): - raise ValueError( - "remove_microbatch_padding=True (sample packing) is not supported for models that " - "pack sequences inside their own forward (e.g. the Qwen3.5 VL Qwen3VLModel): it " - "double-packs and corrupts the GatedDeltaNet cu_seqlens. Set " - "trainer.policy.language_model_only=True to route Qwen3.5 to the native GPTModel GDN " - "packing path, or set trainer.remove_microbatch_padding=False." - ) - # Pending grad-sync request recorded by `_defer_finalize_model_grads`, replayed # by `run_pending_grad_sync`. See those methods for why the sync is deferred. self._pending_grad_sync: Optional[dict] = None @@ -299,12 +290,24 @@ def __call__(self, *args, **kwargs): return self.forward(*args, **kwargs) def _assert_vlm_supported(self): - """Guard the VLM parallelism constraints carried over from the FSDP path. - - 3D RoPE and multimodal token positions make sample/microbatch packing, - context parallelism, and sequence parallelism unsafe for VLMs today. + """Guard the VLM parallelism constraints. + + With ``remove_microbatch_padding`` the model receives a [1, T] THD stream and + Megatron-Bridge's Qwen3-VL model rebuilds 3D mRoPE positions per packed + sub-sequence from ``packed_seq_params`` (``rope.get_rope_index``), without + re-packing the stream (NVIDIA-NeMo/Megatron-Bridge#4532). Other VLMs only run + unpacked (``model_owns_vlm_packing``). + TODO(xgui): context parallelism. preprocess_packed_seqs pre-shards the stream + per CP rank, which the bridge model only accepts with explicit rank-local + 3D position ids. """ - assert not self.remove_microbatch_padding, "VLM + microbatch padding removal unsupported" + if self.remove_microbatch_padding and not self.model_owns_vlm_packing: + model_cls = type(unwrap_model(self.actor_module[0])).__name__ + raise ValueError( + "trainer.remove_microbatch_padding=true (sample packing) is supported for VLMs only on " + "Megatron-Bridge's Qwen3VLModel (Qwen3-VL, Qwen3.5-VL), which rebuilds mRoPE positions per " + f"packed sample; got {model_cls}. Set trainer.remove_microbatch_padding=false." + ) assert mpu.get_context_parallel_world_size() == 1, "VLM + context parallelism unsupported" assert ( mpu.get_tensor_model_parallel_world_size() == 1 or self.cfg.policy.sequence_parallel_size == 1 diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index a66d6b9c3f..0a5f762b26 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -253,8 +253,8 @@ def init_configs( bridge = AutoBridge.from_hf_pretrained(bridge_source, trust_remote_code=True) # For Qwen3.5, language_model_only routes to the native GPTModel + GDN - # path (which supports sample packing) instead of the VL Qwen3VLModel - # (which doesn't). Must run before to_megatron_provider; no-op otherwise. + # path instead of the VL Qwen3VLModel, so the vision tower is not built. + # Must run before to_megatron_provider; no-op otherwise. if language_model_only and maybe_force_qwen35_text_bridge(bridge, hf_config): logger.info( "language_model_only=True: forcing Qwen3.5 text->GPTModel bridge " diff --git a/skyrl/backends/skyrl_train/workers/megatron/model_bridges.py b/skyrl/backends/skyrl_train/workers/megatron/model_bridges.py index 37f9229148..89614dfa1a 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/model_bridges.py +++ b/skyrl/backends/skyrl_train/workers/megatron/model_bridges.py @@ -95,10 +95,9 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM): # Qwen3.5 (language-model-only) -> GPTModel. # - # Qwen3.5 checkpoints dispatch to the VL bridge -> Qwen3VLModel, which packs - # sequences inside its own forward and breaks under SkyRL sample packing. When - # only the LM is wanted (language_model_only=True), route to the native - # GPTModel + GDN thd path instead. The stock text bridges assume a flat text + # Qwen3.5 checkpoints dispatch to the VL bridge -> Qwen3VLModel, which builds + # the vision tower. When only the LM is wanted (language_model_only=True), + # route to the native GPTModel + GDN path instead. The stock text bridges assume a flat text # checkpoint (top-level config, hf_prefix="model."); these subclasses adapt # them to the unified VL checkpoint (text_config, model.language_model.*), # like Qwen35VLBridge but targeting GPTModel. diff --git a/skyrl/train/sft_trainer.py b/skyrl/train/sft_trainer.py index 1df1d28ab9..6f208d0c03 100644 --- a/skyrl/train/sft_trainer.py +++ b/skyrl/train/sft_trainer.py @@ -930,10 +930,12 @@ def setup(self): self.is_vlm = check_is_vlm(self.cfg.trainer.policy.model.path) if self.is_vlm: self.processor = get_processor(self.cfg.trainer.policy.model.path, **tokenizer_kwargs) - # Sequence packing / microbatch padding removal are unsupported for - # VLMs (3D RoPE + image token positions). ``remove_microbatch_padding`` - # defaults to True, so disable both unconditionally and mirror the - # change onto the already-built trainer config the workers receive. + # VLM microbatch padding removal is enabled for RL only (Megatron, on + # Megatron-Bridge's Qwen3VLModel); VLM SFT runs unpacked for now. + # ``remove_microbatch_padding`` defaults to True, so disable both + # unconditionally and mirror the change onto the already-built trainer + # config the workers receive. + # TODO(xgui): enable remove_microbatch_padding for VLM SFT on Megatron. if self.sft_cfg.use_sequence_packing or self.sft_cfg.remove_microbatch_padding: logger.warning("VLM detected: disabling sequence packing / microbatch padding removal.") self.sft_cfg.use_sequence_packing = False diff --git a/tests/backends/skyrl_train/distributed/test_model_owns_vlm_packing.py b/tests/backends/skyrl_train/distributed/test_model_owns_vlm_packing.py new file mode 100644 index 0000000000..0641290560 --- /dev/null +++ b/tests/backends/skyrl_train/distributed/test_model_owns_vlm_packing.py @@ -0,0 +1,35 @@ +"""model_owns_vlm_packing: which VLMs may use sample packing on the Megatron backend.""" + +import pytest +import torch.nn as nn + +megatron_utils = pytest.importorskip("skyrl.backends.skyrl_train.distributed.megatron.megatron_utils") + + +class _OptIn(nn.Module): + model_owns_packing = True + + +def test_plain_model_does_not_own_packing(): + assert not megatron_utils.model_owns_vlm_packing([nn.Linear(2, 2)]) + + +def test_opt_in_attribute(): + assert megatron_utils.model_owns_vlm_packing([_OptIn()]) + assert megatron_utils.model_owns_vlm_packing(_OptIn()) + + +def test_every_chunk_must_own_packing(): + assert not megatron_utils.model_owns_vlm_packing([_OptIn(), nn.Linear(2, 2)]) + + +def test_no_chunks(): + assert not megatron_utils.model_owns_vlm_packing([]) + + +def test_qwen3_vl_model_class(): + model_mod = pytest.importorskip("megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model") + # Skip construction: an uninitialized instance is enough for the isinstance check. + fake = model_mod.Qwen3VLModel.__new__(model_mod.Qwen3VLModel) + nn.Module.__init__(fake) + assert megatron_utils.model_owns_vlm_packing([fake]) diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py index 119b788219..cc13ad9f3c 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py @@ -87,10 +87,9 @@ def get_test_actor_config(model_name) -> SkyRLTrainConfig: cfg.trainer.remove_microbatch_padding = False if "qwen3.5" in model_name.lower(): # Qwen3.5 hybrid GDN checkpoints report a ...ForConditionalGeneration arch - # and auto-dispatch to the VL bridge -> Qwen3VLModel, which self-packs and - # double-packs against SkyRL's sample packing (corrupting the GDN - # cu_seqlens). language_model_only routes them to the native GPTModel + GDN - # thd path instead, which supports packed sequences directly. + # and auto-dispatch to the VL bridge -> Qwen3VLModel. language_model_only + # routes them to the native GPTModel + GDN thd path, which is what this + # text-only test exercises. cfg.trainer.remove_microbatch_padding = True cfg.trainer.policy.language_model_only = True cfg.trainer.ref.language_model_only = True diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py index 80c0cabb8a..386d3f57b8 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_vlm_init.py @@ -26,6 +26,7 @@ from skyrl.train.config.sft_config import SFTConfig, SFTPlacementConfig from skyrl.train.sft_trainer import SFTTrainer from skyrl.train.utils.utils import validate_cfg +from tests.backends.skyrl_train.gpu.gpu_ci.conftest import ray_init from tests.backends.skyrl_train.gpu.utils import ( init_worker_with_type, ray_init_for_tests, @@ -277,12 +278,9 @@ def run_hf_forward(batch, model_name): @pytest.mark.parametrize( ("tp", "cp", "sequence_parallel_size", "remove_microbatch_padding", "gpus_per_node", "expected_substring"), [ - # Qwen3VL packs sequences internally, `remove_microbatch_padding` is not supported - (1, 1, 1, True, 2, "pack sequences inside their own forward"), (2, 1, 2, False, 2, "sequence parallelism"), ], ids=[ - "microbatch_padding", "sequence_parallel", ], ) @@ -362,3 +360,198 @@ async def test_vlm_train(ray_init_fixture, worker_type, tp, pp, gpus_per_node): training_losses.append(step_i_outputs["loss"]) assert training_losses[0] > training_losses[-1] + + +# Packed-vs-unpacked parity. Microbatches of 4 samples (plus a trailing single-sample +# microbatch), so every packed row has samples after the first: a sample-boundary leak +# (wrong mRoPE restart, or GDN state / conv carried across samples) can only show up there. +PACKING_MICRO_BATCH = 4 +PACKING_PROMPTS = [ + ("Describe this picture in one short sentence.", (56, 56), "A small square of noise."), + ("What shapes can you see here? Answer briefly.", (112, 84), "Mostly scattered dots with no clear shape."), + ("Is this image bright or dark?", (224, 168), "It is a mix of bright and dark pixels."), + ("Count the colors.", (84, 224), "Too many colors to count; it looks like random noise."), + ("Give a title for this image.", (140, 140), "Static."), + ("Without any image: what is two plus three?", None, "Two plus three is five."), + ("What would you call this texture?", (196, 112), "A grainy, television-static texture."), + ("Summarize the image.", (70, 154), "Random colored noise in a tall rectangle."), + ("One word for this image?", (168, 56), "Noise."), +] + + +def get_packing_parity_batch(model_name: str) -> TrainingInputBatch: + """Variable-length image (and one text-only) samples laid out like an RL batch. + + Each row is prompt + assistant answer, left-padded. ``response_length`` is the + longest answer and ``loss_mask`` marks each row's answer tokens (right-aligned), + so only text targets are scored -- image-placeholder targets of random-noise + images have huge, bf16-sensitive logprobs and are never trained on. + """ + processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True) + gen = torch.Generator().manual_seed(0) + rows = [] + for prompt, size, answer in PACKING_PROMPTS: + user_content = [{"type": "text", "text": prompt}] + images = None + if size is not None: + from PIL import Image + + w, h = size + pixels = torch.randint(0, 256, (h, w, 3), generator=gen, dtype=torch.uint8).numpy() + images = [Image.fromarray(pixels)] + user_content = [{"type": "image"}] + user_content + messages = [ + {"role": "user", "content": user_content}, + {"role": "assistant", "content": answer}, + ] + text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=False) + prompt_text = processor.apply_chat_template(messages[:1], tokenize=False, add_generation_prompt=True) + out = processor(text=[text], images=images, return_tensors="pt") + prompt_ids = processor(text=[prompt_text], images=images, return_tensors="pt")["input_ids"][0].tolist() + ids = out["input_ids"][0].tolist() + # Score what follows the shared prompt prefix (some templates render the + # generation prompt slightly differently from a completed turn, e.g. think tags). + prefix = next((k for k, (a, b) in enumerate(zip(ids, prompt_ids)) if a != b), len(prompt_ids)) + assert 0 < len(ids) - prefix <= len(ids) // 2, (prefix, len(ids)) + rows.append((ids, len(ids) - prefix, out.get("pixel_values"), out.get("image_grid_thw"))) + + ref_pv = next(pv for _, _, pv, _ in rows if pv is not None) + max_len = max(len(ids) for ids, _, _, _ in rows) + num_actions = max(resp_len for _, resp_len, _, _ in rows) + pad_token_id = processor.tokenizer.pad_token_id or processor.tokenizer.eos_token_id + sequences, attention_mask, loss_mask, pixel_values, image_grid_thw = [], [], [], [], [] + for ids, resp_len, pv, grid in rows: + pad = max_len - len(ids) + sequences.append([pad_token_id] * pad + ids) + attention_mask.append([0] * pad + [1] * len(ids)) + loss_mask.append([0] * (num_actions - resp_len) + [1] * resp_len) + # Text-only rows carry empty vision tensors. + pixel_values.append(pv if pv is not None else ref_pv.new_zeros((0, *ref_pv.shape[1:]))) + image_grid_thw.append(grid if grid is not None else torch.zeros(0, 3, dtype=torch.long)) + + batch_size = len(rows) + loss_mask = torch.tensor(loss_mask, dtype=torch.float) + zeros = torch.zeros(batch_size, num_actions) + data = TrainingInputBatch( + { + "sequences": torch.tensor(sequences), + "attention_mask": torch.tensor(attention_mask), + "action_log_probs": zeros, + "base_action_log_probs": zeros, + "rollout_logprobs": zeros, + "values": zeros, + "returns": zeros, + "advantages": zeros, + "loss_mask": loss_mask, + "response_mask": loss_mask, + "pixel_values": TensorList(pixel_values), + "image_grid_thw": TensorList(image_grid_thw), + } + ) + data.metadata = {"response_length": num_actions} + return data + + +def _megatron_vlm_forward_logprobs( + model_name, batch, tp, remove_microbatch_padding, micro_batch=PACKING_MICRO_BATCH, pp=1 +) -> torch.Tensor: + """Inference forward (the RL old/ref-logprob path): [B, response_length], right-aligned.""" + cfg = get_test_actor_config(model_name=model_name) + cfg.trainer.strategy = "megatron" + cfg.trainer.placement.policy_num_gpus_per_node = tp * pp + cfg.trainer.policy.megatron_config.tensor_model_parallel_size = tp + cfg.trainer.policy.megatron_config.pipeline_model_parallel_size = pp + cfg.trainer.micro_forward_batch_size_per_gpu = micro_batch + cfg.trainer.micro_train_batch_size_per_gpu = micro_batch + cfg.trainer.remove_microbatch_padding = remove_microbatch_padding + # A fresh Ray runtime per forward frees the GPUs for the next one. + with ray_init(): + actor_group = init_worker_with_type( + "policy", shared_pg=None, colocate_all=False, num_gpus_per_node=tp * pp, cfg=cfg + ) + refs = actor_group.async_run_ray_method("mesh", "forward", data=batch) + output = WorkerOutput.cat(actor_group.actor_infos, ray.get(refs)) + return loss_fn_outputs_to_tensor(output.loss_fn_outputs, key="logprobs").float() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model_name", "tp", "pp"), + [ + ("Qwen/Qwen3-VL-2B-Instruct", 1, 1), + ("Qwen/Qwen3-VL-2B-Instruct", 2, 1), + # PP=2: pixel values reach only the first stage; every stage rebuilds mRoPE positions. + ("Qwen/Qwen3-VL-2B-Instruct", 1, 2), + # Qwen3.5 goes through the same Qwen3VLModel but with GatedDeltaNet layers, whose + # state and conv must reset at every packed sample boundary. + ("Qwen/Qwen3.5-0.8B", 1, 1), + ("Qwen/Qwen3.5-0.8B", 2, 1), + ("Qwen/Qwen3.5-0.8B", 1, 2), + ], + ids=["qwen3_vl_tp1", "qwen3_vl_tp2_sp", "qwen3_vl_pp2", "qwen3_5_vl_tp1", "qwen3_5_vl_tp2_sp", "qwen3_5_vl_pp2"], +) +@pytest.mark.megatron +async def test_megatron_vlm_packed_vs_unpacked(model_name, tp, pp): + """Sample packing must add no error beyond ordinary batching. + + Packed and padded (unpacked) microbatches run different kernels on different + shapes, so in bf16 they agree only to rounding -- and so does unpacked with + itself when the microbatch composition changes. The reference is therefore + each sample run alone (unpacked, one sample per microbatch), and the check is + that packed microbatches of 4 are no further from it than unpacked ones. + A sample-boundary leak (mRoPE restart, GatedDeltaNet state/conv carried + across samples) would also make later slots of a packed microbatch worse + than slot 0. + """ + batch = get_packing_parity_batch(model_name) + num_actions = batch.metadata["response_length"] + unpacked = _megatron_vlm_forward_logprobs(model_name, batch, tp, remove_microbatch_padding=False, pp=pp) + packed = _megatron_vlm_forward_logprobs(model_name, batch, tp, remove_microbatch_padding=True, pp=pp) + assert packed.shape == unpacked.shape == (len(PACKING_PROMPTS), num_actions) + + # Reference: every sample alone. At TP>1 (sequence parallel) a microbatch with no image + # crashes Megatron-Bridge's split_deepstack_embs (fixed upstream in Megatron-Bridge#3869, + # not yet merged), so the text-only row is left out of the reference there. + has_image = torch.tensor([size is not None for _, size, _ in PACKING_PROMPTS]) + ref_rows = torch.arange(len(PACKING_PROMPTS)) if tp == 1 else has_image.nonzero().flatten() + ref_batch = TrainingInputBatch({k: (None if v is None else v[ref_rows]) for k, v in batch.items()}) + ref_batch.metadata = batch.metadata + alone = _megatron_vlm_forward_logprobs( + model_name, ref_batch, tp, remove_microbatch_padding=False, micro_batch=1, pp=pp + ) + + scored = batch["loss_mask"].bool() + for name, t in (("unpacked", unpacked), ("packed", packed), ("alone", alone)): + rows_scored = scored[ref_rows] if name == "alone" else scored + assert torch.isfinite(t[rows_scored]).all(), name + ref_scored = scored[ref_rows] + packed_vs_alone = (packed[ref_rows] - alone).abs() + unpacked_vs_alone = (unpacked[ref_rows] - alone).abs() + packed_vs_unpacked = (packed - unpacked).abs() + + slots = (torch.arange(len(PACKING_PROMPTS)) % PACKING_MICRO_BATCH)[ref_rows] + print(f"\n[packing parity] {model_name} tp={tp} pp={pp} (mean/max abs logprob diff on answer tokens)") + for name, d, m in ( + ("packed vs alone", packed_vs_alone, ref_scored), + ("unpacked vs alone", unpacked_vs_alone, ref_scored), + ("packed vs unpacked", packed_vs_unpacked, scored), + ): + print(f" {name:18s}: mean={d[m].mean().item():.5f} max={d[m].max().item():.4f}") + for j, i in enumerate(ref_rows.tolist()): + m = ref_scored[j] + print( + f" sample {i} slot {slots[j].item()}: packed-alone={packed_vs_alone[j][m].mean().item():.5f} " + f"unpacked-alone={unpacked_vs_alone[j][m].mean().item():.5f}" + ) + + packed_err = packed_vs_alone[ref_scored].mean().item() + unpacked_err = unpacked_vs_alone[ref_scored].mean().item() + # Packing must not add error beyond what ordinary batching already shows. + assert packed_err <= 1.5 * unpacked_err + 5e-3, (packed_err, unpacked_err) + # Guard against gross breakage (a wrong position or boundary is >> bf16 noise). + assert packed_vs_unpacked[scored].mean().item() < 5e-2 + # No boundary leak: later slots of a packed microbatch are not worse than slot 0. + first, later = slots == 0, slots > 0 + first_d, later_d = packed_vs_alone[first][ref_scored[first]], packed_vs_alone[later][ref_scored[later]] + assert later_d.mean().item() <= 3 * first_d.mean().item() + 1e-2, (later_d.mean(), first_d.mean()) + assert later_d.max().item() <= 3 * first_d.max().item() + 0.25, (later_d.max(), first_d.max()) diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_packed_vs_unpacked.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_packed_vs_unpacked.py index f78d23919a..9cb8b60113 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_packed_vs_unpacked.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_packed_vs_unpacked.py @@ -33,9 +33,8 @@ from skyrl.train.utils.utils import validate_cfg from tests.backends.skyrl_train.gpu.utils import init_worker_with_type -# MiMo-7B-RL, not Qwen3.5: Qwen3.5's GDN layers cannot sample-pack at all (megatron-LM PR #2644), -# so its recipes run remove_microbatch_padding=false and the wrapper rejects packing outright -- -# there is no packed path to probe. MiMo is dense (Qwen2-style attention) and ships a native MTP +# MiMo-7B-RL, not Qwen3.5: the Qwen3.5 MTP recipes run remove_microbatch_padding=false, so they +# have no packed path to probe. MiMo is dense (Qwen2-style attention) and ships a native MTP # head, and its spec-decode recipe packs for real, so this exercises the path production uses. MODEL_NAME = "XiaomiMiMo/MiMo-7B-RL" @@ -123,7 +122,7 @@ def soft_ce_over_depths(main_logits, students, mask, mask_fn, vp=False): # ========================================================================= # (2) BATCHED LEFT-PADDED, non-packed (remove_left_padding) -- the ACTUAL - # Qwen3.5 training path (GDN can't pack -> REMOVE_MICROBATCH_PADDING=false). + # Qwen3.5 MTP training path (its recipes set REMOVE_MICROBATCH_PADDING=false). # ========================================================================= B = len(token_ids_list) L = max(len(x) for x in token_ids_list) diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_weight_roundtrip.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_weight_roundtrip.py index 6c2b1d9335..1ef22819b5 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_weight_roundtrip.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_mtp_weight_roundtrip.py @@ -109,8 +109,7 @@ def _make_policy_cfg(model_name: str) -> SkyRLTrainConfig: cfg.trainer.policy.megatron_config.tensor_model_parallel_size = 1 cfg.trainer.policy.megatron_config.pipeline_model_parallel_size = 1 cfg.trainer.policy.megatron_config.context_parallel_size = 1 - # This test only round-trips MTP head weights; packing is irrelevant and Qwen3.5's GDN layers - # cannot sample-pack anyway (the wrapper rejects it -- see the 9B recipe's remove_microbatch_padding=false). + # This test only round-trips MTP head weights; packing is irrelevant. cfg.trainer.remove_microbatch_padding = False # Skip optimizer init for weight sync test cfg.trainer.policy.inference_only_init = True