Fix packed compat gpt step arity

Author: intern_nem_dev_1

src/nemotron/recipes/super3/stage1_sft/packed_compat_step.py

@@ -128,8 +128,19 @@ def __getattr__(self, name: str) -> Any:
         return getattr(self._model, name)
 
 
-def forward_step(data_iterator: Any, model: Any):  # pragma: no cover - cluster path
-    """GPT forward step with Mamba packed-sequence keyword compatibility."""
+def forward_step(
+    state_or_data_iterator: Any,
+    data_iterator_or_model: Any,
+    model: Any | None = None,
+    return_schedule_plan: bool = False,
+):  # pragma: no cover - cluster path
+    """GPT forward step with Mamba packed-sequence keyword compatibility.
+
+    Megatron-Bridge has shipped both ``forward_step(data_iterator, model)`` and
+    ``forward_step(state, data_iterator, model, return_schedule_plan=False)``.
+    Keep both call shapes so local static tests and runtime Bridge training use
+    the same compatibility adapter.
+    """
     try:
         from megatron.bridge.training.gpt_step import forward_step as upstream_forward_step
     except ImportError as exc:
@@ -139,8 +150,21 @@ def forward_step(data_iterator: Any, model: Any):  # pragma: no cover - cluster
             "SFT `step_function`."
         ) from exc
 
+    if model is None:
+        data_iterator = state_or_data_iterator
+        model = data_iterator_or_model
+        with _drop_unsupported_packed_seq_params(model) as compat_model:
+            return upstream_forward_step(data_iterator, compat_model)
+
+    state = state_or_data_iterator
+    data_iterator = data_iterator_or_model
     with _drop_unsupported_packed_seq_params(model) as compat_model:
-        return upstream_forward_step(data_iterator, compat_model)
+        return upstream_forward_step(
+            state,
+            data_iterator,
+            compat_model,
+            return_schedule_plan=return_schedule_plan,
+        )
 
 
 __all__ = [

tests/recipes/super3/test_sft_packed_compat_step.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
+import inspect
 import sys
 import types
+from collections.abc import Callable
 from typing import Any
 
 from nemotron.recipes.super3.stage1_sft import packed_compat_step
@@ -47,17 +49,11 @@ def forward(
         return tokens, packed_seq_params
 
 
-def _install_stub_gpt_step() -> None:
+def _install_stub_gpt_step(forward_step: Callable[..., Any]) -> None:
     for parent in ("megatron", "megatron.bridge", "megatron.bridge.training"):
         sys.modules.setdefault(parent, types.ModuleType(parent))
 
     gpt_step = types.ModuleType("megatron.bridge.training.gpt_step")
-
-    def forward_step(data_iterator: Any, model: Any) -> tuple[Any, Any]:
-        del data_iterator
-        output = model(tokens="batch", packed_seq_params={"packed": True})
-        return output, lambda output_tensor: output_tensor
-
     gpt_step.forward_step = forward_step
     sys.modules["megatron.bridge.training.gpt_step"] = gpt_step
 
@@ -72,12 +68,31 @@ def _remove_stub_gpt_step() -> None:
         sys.modules.pop(module_name, None)
 
 
-def test_packed_compat_step_drops_packed_seq_params_for_mamba_like_leaf() -> None:
+def _two_arg_upstream(data_iterator: Any, model: Any) -> tuple[Any, Any]:
+    del data_iterator
+    output = model(tokens="batch", packed_seq_params={"packed": True})
+    return output, lambda output_tensor: output_tensor
+
+
+def test_packed_compat_step_signature_supports_state_aware_bridge_arity() -> None:
+    signature = inspect.signature(packed_compat_step.forward_step)
+
+    assert list(signature.parameters) == [
+        "state_or_data_iterator",
+        "data_iterator_or_model",
+        "model",
+        "return_schedule_plan",
+    ]
+    assert signature.parameters["model"].default is None
+    assert signature.parameters["return_schedule_plan"].default is False
+
+
+def test_packed_compat_step_keeps_existing_two_arg_stub_behavior() -> None:
     leaf = _MambaLikeLeaf()
     model = _ForwardingWrapper(_ForwardingWrapper(leaf))
     assert not packed_compat_step._model_forward_accepts_kwarg(model)
 
-    _install_stub_gpt_step()
+    _install_stub_gpt_step(_two_arg_upstream)
     try:
         output, _loss = packed_compat_step.forward_step(iter(()), model)
     finally:
@@ -87,16 +102,79 @@ def test_packed_compat_step_drops_packed_seq_params_for_mamba_like_leaf() -> Non
     assert leaf.calls == [{"tokens": "batch", "position_ids": None}]
 
 
-def test_packed_compat_step_preserves_packed_seq_params_for_supported_leaf() -> None:
+def test_packed_compat_step_drops_packed_seq_params_for_state_aware_mamba_like_leaf() -> None:
+    upstream_calls: list[dict[str, Any]] = []
+
+    def state_aware_upstream(
+        state: Any,
+        data_iterator: Any,
+        model: Any,
+        return_schedule_plan: bool = False,
+    ) -> tuple[Any, Any]:
+        upstream_calls.append(
+            {
+                "state": state,
+                "data_iterator": data_iterator,
+                "return_schedule_plan": return_schedule_plan,
+            }
+        )
+        output = model(tokens="batch", packed_seq_params={"packed": True})
+        return output, lambda output_tensor: output_tensor
+
+    leaf = _MambaLikeLeaf()
+    model = _ForwardingWrapper(_ForwardingWrapper(leaf))
+    data_iterator = iter(())
+
+    _install_stub_gpt_step(state_aware_upstream)
+    try:
+        output, _loss = packed_compat_step.forward_step("state", data_iterator, model)
+    finally:
+        _remove_stub_gpt_step()
+
+    assert output == ("batch", None)
+    assert leaf.calls == [{"tokens": "batch", "position_ids": None}]
+    assert upstream_calls == [
+        {"state": "state", "data_iterator": data_iterator, "return_schedule_plan": False}
+    ]
+
+
+def test_packed_compat_step_preserves_packed_seq_params_for_state_aware_supported_leaf() -> None:
+    upstream_calls: list[dict[str, Any]] = []
+
+    def state_aware_upstream(
+        state: Any,
+        data_iterator: Any,
+        model: Any,
+        return_schedule_plan: bool = False,
+    ) -> tuple[Any, Any]:
+        upstream_calls.append(
+            {
+                "state": state,
+                "data_iterator": data_iterator,
+                "return_schedule_plan": return_schedule_plan,
+            }
+        )
+        output = model(tokens="batch", packed_seq_params={"packed": True})
+        return output, lambda output_tensor: output_tensor
+
     leaf = _PackedAwareLeaf()
     model = _ForwardingWrapper(_ForwardingWrapper(leaf))
     assert packed_compat_step._model_forward_accepts_kwarg(model)
+    data_iterator = iter(())
 
-    _install_stub_gpt_step()
+    _install_stub_gpt_step(state_aware_upstream)
     try:
-        output, _loss = packed_compat_step.forward_step(iter(()), model)
+        output, _loss = packed_compat_step.forward_step(
+            "state",
+            data_iterator,
+            model,
+            return_schedule_plan=True,
+        )
     finally:
         _remove_stub_gpt_step()
 
     assert output == ("batch", {"packed": True})
     assert leaf.calls == [{"tokens": "batch", "packed_seq_params": {"packed": True}}]
+    assert upstream_calls == [
+        {"state": "state", "data_iterator": data_iterator, "return_schedule_plan": True}
+    ]

workspace/interns/intern_nem_dev_1/status.md

@@ -1,12 +1,12 @@
 # intern_nem_dev_1 - Status
 
-<!-- METADATA:STATUS=Working,TASK=task211_qwen_sft_mamba_packed_seq_params_compat_s1,ROLE=dev,SESSION=2 -->
+<!-- METADATA:STATUS=Working,TASK=task213_qwen_sft_packed_compat_gpt_step_arity_s1,ROLE=dev,SESSION=1 -->
 
 | Field | Value |
 |------|-----|
 | Name | intern_nem_dev_1 |
 | Status | Working |
-| Current Task | task211_qwen_sft_mamba_packed_seq_params_compat_s1 |
-| PR | https://github.com/songCNMS/Nemotron/pull/309 |
-| Session | 2 |
-| Recent Progress | PR #309 open and mergeable; final checks passed; pushing Session 2 status/report closeout |
+| Current Task | task213_qwen_sft_packed_compat_gpt_step_arity_s1 |
+| PR | Pending |
+| Session | 1 |
+| Recent Progress | Accepted task213; fixed packed compat adapter state-aware gpt_step arity and started validation |

workspace/tasks/task213_qwen_sft_packed_compat_gpt_step_arity_s1/README.md

@@ -0,0 +1,45 @@
+# task213_qwen_sft_packed_compat_gpt_step_arity_s1
+
+<!-- METADATA:STATUS=InProgress,ASSIGNEE=intern_nem_dev_1,SESSION=1 -->
+
+## Scope
+
+- Fix the task212 evidence failure in
+  `src/nemotron/recipes/super3/stage1_sft/packed_compat_step.py`.
+- Preserve Megatron-Bridge state-aware `gpt_step` arity:
+  `(state, data_iterator, model, return_schedule_plan=False)`.
+- Keep compatibility with existing local two-argument test stubs where
+  reasonable.
+- Continue filtering `packed_seq_params` only around the model argument, and
+  only for model forward chains that do not support that keyword.
+
+## Boundaries
+
+- Product code/tests/docs only.
+- No live train, package install, endpoint, benchmark, W&B, cluster/deploy,
+  artifact upload, direct `main`/`master` push, or self-merge.
+
+## Status
+
+- Base: `f65dafdb15b28342c1fbd4a5ead807052bcdd264`.
+- Branch: `intern_nem_dev_1/task213_qwen_sft_packed_compat_gpt_step_arity_s1`.
+- PR: pending.
+- Current implementation:
+  - `packed_compat_step.forward_step` now accepts the runtime
+    state-aware Bridge call shape and passes `state`, `data_iterator`, `model`,
+    and `return_schedule_plan` through to upstream `gpt_step`.
+  - Existing two-argument stub behavior is retained for local static tests.
+  - Focused tests now cover state-aware Mamba-like drop, packed-aware preserve,
+    `return_schedule_plan` propagation, legacy two-arg behavior, and dispatch
+    wiring.
+- Checks completed so far:
+  - `PYTHONPATH=src /work-agents/.venv/bin/python -m pytest -q tests/recipes/super3/test_sft_packed_compat_step.py tests/recipes/super3/test_sft_forward_step_dispatch.py`
+    -> `10 passed, 1 skipped`.
+  - `PYTHONPATH=src /work-agents/.venv/bin/python -m pytest -q tests/recipes/super3/test_m1_agentic_sft.py -k 'qwen_local_train or qwen30b_a3b_local_train'`
+    -> `10 passed, 86 deselected`.
+  - `PYTHONPATH=src /work-agents/.venv/bin/python -m pytest -q tests/recipes/super3/test_qwen_chat_contract.py tests/recipes/super3/test_stage1_sft_default_config.py tests/recipes/super3/test_m1_agentic_sft.py -k 'qwen or sft or packed or data_prep or target_family'`
+    -> `136 passed, 1 skipped`.
+  - `/work-agents/.venv/bin/ruff check src/nemotron/recipes/super3/stage1_sft/packed_compat_step.py tests/recipes/super3/test_sft_packed_compat_step.py tests/recipes/super3/test_sft_forward_step_dispatch.py`
+    -> passed.
+- Blockers: none currently.
+- Residual risk: no live SFT training rerun per PM boundary.

workspace/tasks/task213_qwen_sft_packed_compat_gpt_step_arity_s1/history_log.md

@@ -0,0 +1,20 @@
+# History Log
+
+<!-- METADATA:SESSION=1 -->
+
+## Session 1 - 2026-05-30
+
+- Accepted PM task213 on branch
+  `intern_nem_dev_1/task213_qwen_sft_packed_compat_gpt_step_arity_s1`
+  from base `f65dafdb15b28342c1fbd4a5ead807052bcdd264`.
+- Root cause from task212 evidence: task211 compatibility adapter called
+  upstream `gpt_step.forward_step(data_iterator, compat_model)`, but runtime
+  Megatron-Bridge uses state-aware signature
+  `(state, data_iterator, model, return_schedule_plan=False)`.
+- Updated `packed_compat_step.forward_step` to support the state-aware call
+  shape while retaining two-argument local stub compatibility.
+- Added focused tests for state-aware Mamba-like filtering, packed-aware
+  preservation, `return_schedule_plan` propagation, legacy two-arg behavior,
+  and adapter signature/dispatch wiring.
+- Began focused and broader SFT/Qwen validation; no live training or forbidden
+  operations were run.

workspace/tasks/task213_qwen_sft_packed_compat_gpt_step_arity_s1/task_knowledge.md

@@ -0,0 +1,14 @@
+# Task Knowledge
+
+<!-- METADATA:SESSION=1 -->
+
+- Runtime Megatron-Bridge `gpt_step.forward_step` in the task212 environment
+  has signature `(state, data_iterator, model, return_schedule_plan=False)`.
+- The task211 adapter failed because it only called upstream
+  `forward_step(data_iterator, compat_model)`, producing
+  `TypeError: forward_step() missing 1 required positional argument: 'model'`.
+- `packed_compat_step.forward_step` must wrap only the model argument with
+  `_drop_unsupported_packed_seq_params`; `state`, `data_iterator`, and
+  `return_schedule_plan` must pass through unchanged.
+- The local two-argument test stub path is still useful because it keeps
+  sandbox tests independent of an installed Megatron-Bridge runtime.