[dyurk-lila]

What breaks

On the Megatron backend, enabling both transformer_config_kwargs.fine_grained_activation_offloading: true and trainer.fused_lm_head_logprob: true crashes in the policy forward:

AttributeError: type object 'FineGrainedActivationOffloadingInterface' has no attribute 'mark_not_offloadable'. Did you mean: 'mark_not_offload'?

Root cause

call_model_with_fused_lm_head (skyrl/backends/skyrl_train/distributed/megatron/fused_lm_head.py) turns off Megatron post-processing and handles the LM head itself. Because of that, it has to mark the output-layer parameters as not offloadable on its own, and it does so with FineGrainedActivationOffloadingInterface.mark_not_offloadable. That method doesn't exist at the pinned megatron-core commit (07147d6942dfd4e6d1566b64bb1730f478a1bc32). The real method is the staticmethod mark_not_offload(tensor) in megatron/core/pipeline_parallel/fine_grained_activation_offload.py. The unit test's stub of the interface used the same wrong name, so the tests never caught the bug.

Fix

• Call off_interface.mark_not_offload(parameter). This is the only call site.
• Rename the stub in tests/backends/skyrl_train/distributed/test_fused_lm_head.py to mark_not_offload, and have it set the same _do_not_offload flag that Megatron's implementation sets.
• Add test_megatron_offload_interface_exposes_mark_not_offload. It checks that the installed megatron-core interface has mark_not_offload, and skips if megatron-core isn't installed. If the method is renamed again, this test fails, so the stub can't silently drift from the real interface.

Tests

• At the pinned megatron-core: hasattr(FineGrainedActivationOffloadingInterface, "mark_not_offload") is True, and hasattr(..., "mark_not_offloadable") is False.
• pytest tests/backends/skyrl_train/distributed/test_fused_lm_head.py: 14 passed (CPU).
• With the corrected stub and the original source line, the hybrid test fails with the same AttributeError shown above.
• ruff (v0.11.9) and black (24.10.0) are clean on the changed files.

🤖 Generated with Claude Code

---

Note

Low Risk
Single-call-site rename to match megatron-core; fixes a runtime AttributeError with no auth or data-path changes.

Overview
Fixes a crash when fine-grained activation offloading and fused LM head logprob are both enabled on the Megatron HybridModel path: call_model_with_fused_lm_head now calls Megatron’s FineGrainedActivationOffloadingInterface.mark_not_offload instead of the non-existent mark_not_offloadable when marking output-layer parameters before bypassing post-processing.

Tests are aligned with the real API (_do_not_offload on tensors), and a new integration-style check asserts that installed megatron-core exposes mark_not_offload so the stub cannot drift again.

^{Reviewed by Cursor Bugbot for commit 00a5f96. Bugbot is set up for automated code reviews on this repo. Configure here.}

[gemini-code-assist]

Code Review

This pull request updates the activation offloading interface from mark_not_offloadable to mark_not_offload to align with changes in the Megatron API, and updates the corresponding tests. A review comment suggests improving the new guard test by skipping only if the top-level megatron package is missing, rather than the deep submodule path, to prevent the test from being silently skipped if the submodule is renamed or moved.

[gemini-code-assist]

Using pytest.importorskip on the deep submodule path megatron.core.pipeline_parallel.fine_grained_activation_offload will cause the test to be silently skipped if the submodule is renamed, moved, or removed in future Megatron versions. This defeats the purpose of this guard test, which is to detect drift or changes in the Megatron interface. Instead, skip only if the top-level megatron package is missing, and perform a standard import for the submodule so that any internal changes or drift will trigger a test failure.

Suggested change

	offload = pytest.importorskip("megatron.core.pipeline_parallel.fine_grained_activation_offload")
	pytest.importorskip("megatron")
	import megatron.core.pipeline_parallel.fine_grained_activation_offload as offload

[greptile-apps]

Confidence Score: 4/5

[Medium impact] The change reaches a bounded Megatron training path.

The fix appears safe to merge, with a non-blocking gap in the new test's CI coverage.

Findings

1.  CI skips the interface check ▶

_{Reviews (1) · Last reviewed commit: "[fix][megatron] Call mark_not_offload in..." · Reviewed by Greptile}

[greptile-apps]

CI skips the interface check

The new test never checks the real Megatron interface in CPU CI. The job that installs Megatron selects only tests marked megatron, but this test has no marker. The other jobs lack Megatron, so pytest.importorskip skips it. Add @pytest.mark.megatron to this test so CI can catch another method-name change instead of letting the stub tests pass while training breaks.

Suggested change

	def test_megatron_offload_interface_exposes_mark_not_offload():
	@pytest.mark.megatron
	def test_megatron_offload_interface_exposes_mark_not_offload():