feat(engine): nemo_automodel engine for unified engine workers (SFT + PPO/GRPO) - #1
Draft
HuiyingLi wants to merge 9 commits into
Draft
feat(engine): nemo_automodel engine for unified engine workers (SFT + PPO/GRPO)#1HuiyingLi wants to merge 9 commits into
HuiyingLi wants to merge 9 commits into
Conversation
…e workers Adds a model-engine backend that delegates training to nemo_automodel's in-process Engine (forward/extraction/microbatch-lifecycle/optimizer), so the unified engine workers (SFT + PPO/GRPO) can use Automodel as the actor/ref. - verl/workers/engine/automodel/engine_tinker_impl.py: AutomodelEngine (BaseEngine), builds the model via the current nemo_automodel API (build_model + create_distributed_setup_from_config) and delegates the step to nemo_automodel.engine.Engine via its PackedBatch pass-through. Bridges the Engine's per-datum ModelOutput to verl's flat model_output (log_probs + entropy). Implements infer_batch (compute_log_prob), get_per_tensor_param (HF-named rollout weight sync), save/load_checkpoint, disable_adapter. - AutomodelActorConfig + config groups (actor/ref/critic, model_engine) so the PPO/GRPO path can select `model_engine=automodel` (previously automodel was SFT-only). Validated end-to-end: 8-GPU GRPO on gsm8k (Qwen2.5-0.5B) — rollout (vLLM) -> compute_old_log_prob -> GRPO advantage -> actor update -> weight sync, two steps with full metrics. Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
…gine.py move The Automodel engine loads a custom Qwen2 forward whose rotary defaults to fused TE RoPE (BackendConfig.rope_fusion = HAVE_TE and cuda). The fused kernel indexes rotary angles by physical sequence position and assumes contiguous [0, seq_len), so it ignores the per-sequence position_id resets of a THD-packed micro-batch and corrupts RoPE for every non-first sequence in a pack. In GRPO this made the actor's recomputed logprobs drift from the vLLM rollout (rollout_actor_probs_pearson_corr ~0.62 vs fsdp's 0.9997), breaking the importance ratio so the policy didn't learn. Build the model with backend rope_fusion=False so the non-fused path gathers cos/sin by position_id value (packing-correct). corr -> 0.9997, matching fsdp. Also update the import to follow nemo_automodel.engine -> components.training.engine. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
train_mode/eval_mode previously only flipped module train()/eval() and skipped
offload ("offload omitted for now"), so with param_offload/optimizer_offload
enabled the actor was pushed to CPU after rollout weight-sync but never onloaded
back before the next step — breaking the actor<->vLLM swap for large models that
must offload to fit the rollout engine.
Replace the custom _ModeCtx with EngineTrainModeCtx/EngineEvalModeCtx subclassing
BaseEngineCtx, exactly like the fsdp/megatron/torchtitan/veomni backends:
BaseEngineCtx._context_switch drives onload-on-enter / offload-on-exit (gated by
is_param_offload_enabled / is_optimizer_offload_enabled, honoring
disable_auto_offload) via the engine's to(), and the subclass only adds the
per-part train()/eval() toggle (+ zero-grad before offload on train exit).
Verified single-GPU round-trip with offload enabled: params/optimizer move
cuda->cpu on context exit and back cuda on the next enter. With offload disabled
(default) _context_switch is a no-op, preserving the co-resident path.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds an automodel model-engine backend that delegates training to
nemo_automodel's in-process
Engine, so the unified engine workers can use Automodel as the actor/ref forboth SFT and PPO/GRPO (previously the automodel engine was SFT-only).
What's added
verl/workers/engine/automodel/engine_tinker_impl.py—AutomodelEngine(BaseEngine):build_model+create_distributed_setup_from_config) and wraps it innemo_automodel.engine.Engine;forward_backward_batchruns verl's THD micro-batches through the Engine'sPackedBatchpass-through (multi-microbatch grad accumulation), bridging theEngine's per-datum
ModelOutputback to verl's flatmodel_output["log_probs"]/["entropy"]so the existingppo_loss/sft_lossrun unchanged;infer_batch(compute_log_prob),get_per_tensor_param(HF-named rolloutweight sync),
save/load_checkpoint,disable_adapter.AutomodelActorConfig+ config groups (actor/automodel_actor,ref/automodel_ref,critic/automodel_critic,model_engine/automodel) forthe
model_engine=automodelPPO/GRPO path.Usage
Validation
End-to-end 8-GPU GRPO on gsm8k (Qwen2.5-0.5B): rollout (vLLM) →
compute_old_log_prob→ GRPO advantage → actor update → vLLM weight sync, twotraining steps with full metrics (
actor/pg_loss,actor/grad_norm,critic/rewards,timing_s/update_weights). Also a single-GPU functionalforward/backward step.
Notes
automodel/transformer_impl.py(pinned to an older Automodel API)is superseded by this engine;
__init__now registers the new one. Removingthe legacy file is left as a follow-up.
scaling in extraction.
🤖 Generated with Claude Code