[yueming-yuan]

Co-authored with @yushengsu-thu.

Expose command-level training on MultiLoRATrainRayActor in megatron_utils/lora/actor.py: forward_backward accumulates gradients across calls; optim_step(adam_params_by_slot) steps only the requested slots; forward_only_logprobs scores tokens; and load_slot / unload_slot manage slot weights, gradients, and optimizer state.

lora/model.py extends the shared model execution through run_forward_backward, which shares run_forward_backward_pass and setup_train_iteration_config with the existing training path. Adam parameters are applied per request. Successful steps and nonfinite-gradient skips return per-slot outcomes; unknown preparation, optimizer, and collective failures propagate to the existing TrainerCell failure boundary. Slot optimizer construction isolates the base optimizer factory while preserving LayerWise wrapping and FP32 state. The command dispatcher requires one live trainer cell. Existing optimizer tests cover nonfinite skips and propagation of unknown execution failures.

Slot optimizer construction passes the model chunks to LayerWiseDistributedOptimizer so equal-size adapter parameters use Megatron's rank-independent ownership tie-breaker.

Multi-LoRA pool initialization returns the model with no pool optimizer or scheduler. Only per-slot optimizers own training state; scheduler construction is skipped at this layer.

Stacked on #2839.

The multi-LoRA actor inherits the common initialization lifecycle and owns the slot optimizer registry. MegatronTrainRayActor._init_training_state() initializes the ordinary weight updater; the multi-LoRA override initializes slot state. The worker spec selects the subclass only for Megatron multi-LoRA actor workers. Full-model, single-LoRA, FSDP, and critic selection remain unchanged.

Validation:

• pre-commit run --all-files and the incremental git diff --check passed on this component head.
• Cascade verification checks that moved methods retain their behavior and that the final trees match the reviewed code. Formatting is committed separately.
• No new test files were added. The existing worker-class selection test now covers the multi-LoRA actor and preserves the critic's backend selection.
• Fresh CPU/GPU CI is pending. Combined validation and limitations are recorded on [DO NOT MERGE] stacked multi-lora full PR #2846 and Unify HF snapshot and checkpoint writing #3198.

ci-megatron-pr: 73b54618f7e58e0f25f619bcfecbe2640765475a
ci-sglang-pr: 02b5e123c490c988d21526d0dd7837fefdd0e1da

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}