[Feature]: Lora support

[princepride]

Summary

#206
This PR adds end-to-end LoRA support for Megatron actor training with colocated vLLM rollout. When --lora-rank > 0, the training actor applies LoRA adapters via Megatron-Bridge, exports PEFT-format adapters after each weight update, and hot-loads them into vLLM through the runtime LoRA API. Rollout requests then use the adapter name instead of the base model path.

This enables parameter-efficient RL fine-tuning without full-weight sync between trainer and inference engines on every step.

Motivation

Full-weight synchronization between Megatron and colocated vLLM engines is expensive for large models. LoRA reduces the sync payload to small adapter weights and lets vLLM load adapters at runtime, which is a better fit for iterative RL training.

Key Changes

Training (Megatron)

• New CLI args: --lora-rank, --lora-alpha, --lora-dropout, --lora-type (lora / canonical_lora), --target-modules, --exclude-modules, --lora-adapter-name
• lora_utils.py: target module parsing/conversion (HF ↔ Megatron), PEFT config generation, adapter export (adapter_model.bin + adapter_config.json)
• model_provider.py: apply LoRA directly after Bridge provider returns the actor model (instead of relying on pre-wrap hooks)
• actor.py: wake up offloaded model before LoRA export in offload_train + colocate mode
• arguments.py: auto-enable vllm_enable_lora, validate LoRA requires --megatron-to-hf-mode bridge and --colocate

Weight Sync (Trainer → vLLM)

• update_weight_from_tensor.py: when LoRA is enabled, skip full-weight IPC sync; export adapter and call load_lora_adapter on all rollout engines
• vllm_engine.py:
  • set VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 when LoRA is enabled
  • add load_lora_adapter() with unload-then-load flow
  • handle vLLM's non-JSON / empty 200 responses on LoRA endpoints

Rollout

• vllm_rollout.py / vllm_streaming_rollout.py: use adapter name (default vime_lora) as the "model" field in rollout requests when LoRA is enabled

Cleanup & Reliability

• train.py: wrap main loop in try/finally to dispose actor/rollout resources
• actor_group.py / actor.py: add dispose() for explicit wandb cleanup in Ray actors
• logging_utils.py: call wandb.teardown() after wandb.finish() to avoid BrokenPipeError noise on Ray actor exit

Tests

python3 train.py \
  --hf-checkpoint /root/models/Qwen3-8B \
  --ref-load /root/models/Qwen3-8B_torch_dist \
  --load /root/models/Qwen3-8B \
  --finetune \
  --no-load-optim \
  --no-load-rng \
  --save /root/vime_lora_runs/qwen3-8b-dapo-lora-smoke \
  --save-interval 1 \
  --megatron-to-hf-mode bridge \
  --actor-num-nodes 1 \
  --actor-num-gpus-per-node 1 \
  --colocate \
  --calculate-per-token-loss \
  --prompt-data /root/datasets/dapo-math-17k/dapo-math-17k.jsonl \
  --input-key prompt \
  --label-key label \
  --apply-chat-template \
  --rollout-shuffle \
  --rm-type deepscaler \
  --num-rollout 1 \
  --rollout-batch-size 1 \
  --n-samples-per-prompt 1 \
  --rollout-max-response-len 1024 \
  --rollout-temperature 0.8 \
  --global-batch-size 1 \
  --micro-batch-size 1 \
  --advantage-estimator grpo \
  --kl-loss-coef 0.00 \
  --kl-loss-type k1 \
  --kl-coef 0.00 \
  --entropy-coef 0.00 \
  --eps-clip 4e-4 \
  --optimizer adam \
  --lr 1e-6 \
  --lr-decay-style constant \
  --weight-decay 0.1 \
  --adam-beta1 0.9 \
  --adam-beta2 0.98 \
  --tensor-model-parallel-size 1 \
  --pipeline-model-parallel-size 1 \
  --context-parallel-size 1 \
  --expert-model-parallel-size 1 \
  --expert-tensor-parallel-size 1 \
  --use-dynamic-batch-size \
  --max-tokens-per-gpu 4096 \
  --rollout-num-gpus-per-engine 1 \
  --rollout-num-gpus 1 \
  --vllm-gpu-memory-utilization 0.45 \
  --vllm-max-cudagraph-capture-size 32 \
  --lora-rank 8 \
  --lora-alpha 16 \
  --lora-dropout 0.0 \
  --lora-type lora \
  --target-modules all-linear \
  --only-train-params-name-list lora_A lora_B linear_in linear_out \
  --lora-adapter-name vime_lora \
  --attention-dropout 0.0 \
  --hidden-dropout 0.0 \
  --accumulate-allreduce-grads-in-fp32 \
  --attention-softmax-in-fp32 \
  --attention-backend flash \
  --use-wandb \
  --wandb-host https://api.wandb.ai \
  --wandb-project vime-lora-smoke \
  --wandb-group qwen3-8b-dapo-single-h100 \
  --wandb-key "${WANDB_API_KEY}" \
  --disable-wandb-random-suffix

---

Supersedes #228 (re-opened from fork branch princepride:lora-support).

[read-the-docs-community]

Documentation build overview

📚 vime | 🛠️ Build #33532168 | 📁 Comparing 1350250 against latest (2702eef)

  🔍 Preview build  

42 files changed · ± 41 modified · - 1 deleted

± Modified

• genindex.html
• index.html
• search.html
• advanced/arch-support-beyond-megatron.html
• advanced/delta-weight-sync.html
• advanced/external-rollout-engines.html
• advanced/fault-tolerance.html
• advanced/low-precision.html
• advanced/megatron-config.html
• advanced/observability.html
• and 31 more...

- Deleted

• platform_support/amd_tutorial.html

[gemini-code-assist]

Code Review

This pull request introduces support for actor LoRA training and runtime vLLM adapter serving. It adds new command-line arguments, integrates PEFT LoRA adapter saving and loading into the checkpointing mechanism, and updates the weight update flow to support hot-reloading adapters on vLLM rollout engines. The review feedback identifies two key issues: a potential deadlock during weight updates if Ray calls on rank 0 fail, which can be resolved by broadcasting a status flag to other ranks, and a failure to respect the 'skip_load_to_model_and_opt' flag when loading LoRA adapters during checkpoint resume.

[gemini-code-assist]

If any of the Ray calls (like pause_generation or flush_cache) on rank 0 raise an exception, rank 0 will enter the finally block and call dist.barrier(). Ranks 1..N will also reach their dist.barrier() and unblock. However, rank 0 will then propagate the exception and crash/exit, while ranks 1..N will proceed to save_lora_adapter_for_vllm(). Inside save_lora_adapter_for_vllm(), ranks 1..N will call dist.barrier(). Since rank 0 has already crashed and will never reach this barrier, all other ranks will hang indefinitely.

To prevent this, we should broadcast a success/failure status flag from rank 0 to all other ranks using dist.broadcast on a CPU tensor, and raise an exception on other ranks if rank 0 failed.

    @torch.no_grad()
    def _update_lora_adapter(self) -> None:
        """Export the current adapter and ask vLLM engines to load it by path."""
        rank = dist.get_rank()
        status = torch.tensor([0], dtype=torch.int)
        if rank == 0:
            try:
                ray.get([engine.pause_generation.remote() for engine in self.rollout_engines])
                ray.get([engine.flush_cache.remote() for engine in self.rollout_engines])
                status[0] = 1
            except Exception as e:
                logger.error("Failed to pause generation or flush cache on rollout engines: %s", e)
                raise e
            finally:
                dist.broadcast(status, src=0, group=get_gloo_group())
        else:
            dist.broadcast(status, src=0, group=get_gloo_group())
            if status[0] == 0:
                raise RuntimeError("Rank 0 failed to pause generation or flush cache on rollout engines.")

        adapter_path = save_lora_adapter_for_vllm(self.model, self.args, self.weight_version)

        if rank == 0:
            try:
                refs = [
                    engine.load_lora_adapter.remote(
                        self.args.lora_adapter_name,
                        adapter_path,
                        weight_version=str(self.weight_version),
                    )
                    for engine in self.rollout_engines
                ]
                ray.get(refs)
                ray.get([engine.continue_generation.remote() for engine in self.rollout_engines])
                status[0] = 1
            except Exception as e:
                status[0] = 0
                logger.error("Failed to load LoRA adapter or continue generation on rollout engines: %s", e)
                raise e
            finally:
                dist.broadcast(status, src=0, group=get_gloo_group())
        else:
            dist.broadcast(status, src=0, group=get_gloo_group())
            if status[0] == 0:
                raise RuntimeError("Rank 0 failed to load LoRA adapter or continue generation on rollout engines.")

[gemini-code-assist]

When skip_load_to_model_and_opt is True, we are supposed to skip loading weights to the model and optimizer. However, load_lora_adapter is currently called unconditionally, which will load the adapter weights into the model and optimizer state.

We should respect skip_load_to_model_and_opt by passing it to load_lora_adapter as skip_load_to_model, and only passing optimizer and opt_param_scheduler if skip_load_to_model_and_opt is False.

    if is_lora_enabled(args):
        adapter_path = getattr(args, 

[aoshen02]

We don't add more unit test here.