[multi-lora] 1/n operation backend: explicit training with Tinker compatibility

[yushengsu-thu]

Summary

Replaces the old dataset-driven Multi-LoRA path with a service-only, client-driven operation backend. Multiple clients share one Megatron base model through fixed LoRA slots and explicitly control forward/backward, optimizer, checkpoint, and sampler-publication boundaries.

This PR provides the backend and Ray-side operation contract. Stacked PR #2346 adds the official tinker==0.24.1 JSON/REST frontend, and #2365 adds the complete client-owned SFT/RL guide.

Reviewer

Tinker SDK client
    │
    │ JSON/REST translation (#2346; not part of this PR)
    ▼
TinkerFrontendHTTPServer (#2346)
    │ calls the backend inside the head-node controller actor
    ▼
MultiLoraOperationController
    └── MultiLoraOperationBackend
        ├── AdapterRegistry
        │   └── fixed slot for (name, registration_id)
        ├── OperationLedger
        │   └── ordering, retries, gaps, ACK/backpressure, and fencing
        ├── FixedSlotResidency
        │   └── claim-and-bind → immutable BatchExecutionLease
        ├── GradientWindowTracker
        └── RouterInferenceAdmin

main()
│
├─ 138  actor_model, _ = create_training_models(...)
│        └─ actor_model = v1 RayTrainGroup
│           └─ _actor_handles = [MegatronTrainRayActor handle × one per GPU rank]
│
└─ while True
   │
   ├─ A. Adapter residency
   │  actor_model.reconcile_tinker_adapters()
   │  └─ RayTrainGroup._broadcast("reconcile_tinker_adapters")
   │     └─ Each MegatronTrainRayActor.reconcile_tinker_adapters()
   │        └─ multi_lora.trainer.reconcile_adapters(...)
   │           └─ Read the controller registry snapshot, then load or unload LoRA slots
   │
   ├─ B. Control operations
   │  run_control_phase(...)
   │  ├─ controller.claim_ready_control_operations()
   │  ├─ actor_model.execute_tinker_controls(operations, lease)
   │  │  └─ Broadcast to every MegatronTrainRayActor
   │  │     └─ multi_lora.trainer.execute_controls(...)
   │  └─ weight_updater.update_weights()
   │     └─ actor_model.update_weights()
   │        └─ Publish pending LoRA adapter weights to serving
   │
   └─ C. Data training
      ├─ inference_controller.prepare_rollout(rollout_id)
      ├─ rollout_executor.generate(rollout_id)
      │  └─ Produce this batch of Tinker rollout_data
      └─ train_data_batch(actor_model, controller, rollout_id, rollout_data)
         └─ actor_model.train(rollout_id, rollout_data)
            └─ RayTrainGroup._broadcast("train", ...)
               └─ Each MegatronTrainRayActor.train(...)
                  ├─ Run a Megatron training step
                  └─ On success:
                     multi_lora.trainer.commit_batch(...)
                     └─ controller.commit_tinker_batch(...)

MultiLoraParameterExecutor is deliberately narrower than the whole backend: it owns optimizer step/discard for lease-bound adapter slots. Data conversion and forward/backward live in MultiLoraOperationBatchFn; checkpoint and publication remain Multi-LoRA trainer controls. create_rollout_components() also exposes separate InferenceControllerPort and RolloutExecutorPort roles. Today they are two adapters over the same combined RolloutManager; the physical controller/executor split remains future integration work.

Operation semantics

Supported operations are:

• forward
• forward_backward
• optim_step
• save_weights_for_sampler
• save_state
• load_state

Important invariants:

• Operations execute in strict ordinal order per registration, while out-of-order arrival is gap-buffered.
• Identical retries are idempotent; reused identities with different content are conflicts.
• (name, registration_id) prevents stale handles from targeting a re-registered adapter with the same display name.
• Claim-and-bind produces one immutable execution lease, which trainer ranks validate before physical mutation.
• forward_backward calls accumulate one client-owned gradient window. A failed batch poisons that window, and the next optim_step discards it instead of applying partial gradients.
• Results remain available until ACK; capacity pressure rejects new work instead of evicting terminal results.
• Sampler publication completes only after staged weights are live.

Sampling, scoring, rewards, advantages, batch scheduling, and SDK Datum
construction remain client-owned.

Full-parameter seam

The reusable seam is intentionally narrower than the complete backend:

• BatchExecutionLease
• TrainerResidencyPort
• ParameterExecutor
• run_optim_controls

FullParameterExecutor implements singleton whole-model Adam step/discard
behavior against a stock Megatron optimizer. It is not connected to launch
configuration, registration, data conversion, forward/backward execution,
checkpointing, controller routing, or sampler publication. This PR therefore
does not claim working full-parameter SFT/RL or full-parameter GPU/E2E support.

Validation

Current head 7189b1e54 has:

• 29 successful GitHub checks, 2 expected skips, and no failures or
  pending checks.
• CPU test shards, pre-commit, and CodeQL.
• 2/4/8-GPU H200, 8-GPU H100, and 4-GPU MI350 CI coverage.
• Development-time 2×H200 acceptance covering Qwen SFT/RL, GPT-OSS
  multi-adapter training/checkpointing, sampler publication, restore,
  registration fencing, and slot reuse.

Limits

• Text-only, 1-D shifted targets with cross_entropy,
  importance_sampling, or ppo.
• Fixed adapter slots; no eviction, idle-slot GC, or per-tenant quota.
• Megatron backend with Adam-family per-slot optimization.
• Disaggregated execution only; pipeline parallelism must be 1 and
  qkv_format must be thd.
• Latest-only sampler publication; no immutable version-pinned snapshots.
• Restore rejects incompatible world topology, LoRA shape, or per-rank
  optimizer ownership.
• Recovery of work lost after claim during downstream process death is not
  guaranteed; executor-side reconciliation remains future work.

Stack and dependencies

• [multi-lora] 2/n Tinker frontend: tinker==0.24.1 JSON/REST adapter #2346 — tinker==0.24.1 JSON/REST frontend.
• [multi-lora] 3/n Tinker full stack: backend, frontend, and RL/SFT docs #2365 — full-stack SFT/RL examples and documentation.
• [optimizer] Fix LayerWise expert ownership Megatron-LM#82 — LayerWise expert ownership; merged into
  miles-main.
• peft: fix multi-LoRA activation-recompute replay and grouped-expert adapter DDP flagging Megatron-Bridge#27 — Multi-LoRA recompute and grouped-expert
  support; merged into the tracked bridge branch.

0821 liveness fixes: capped tolerance for consecutive generate failures (--multi-lora-max-consecutive-generate-failures), a claimed-operation TTL backstop that unblocks orphaned CLAIMED queue heads (--tinker-operation-claimed-ttl), and FAILED child runtimes self-healing to IDLE after a cooldown.

• 0821: fixed the publish-path function-local import stranded by the api_backends regrouping (update_weight .../mixin.py → api_backends.multi_lora.model) + tests/fast/test_import_integrity.py static import-integrity regression; 2-GPU mini-loop re-verified publish+sample (PASS).

• 0822 test cleanup: launch-unreachable multi-LoRA tests removed (stamped-slot fallback deleted from production, now a loud ValueError), duplicate trim test dropped, sglang abort regression moved to its own PR [multi-lora] 1/n - 7, test: pin update_weight_version abort_all_requests=False regression #2713, RL-quality harness now exits non-zero on aborted loops.

[yushengsu-thu]

clean

[yushengsu-thu]

@claude review

[yushengsu-thu]

@claude review

[claude]

I reviewed this update to the multi-LoRA operations PR and found no new bugs this run. Previously-flagged issue in this bot's own prior review — the OperationLedger memory leak for retired registrations — is now fixed: fence() (operations.py:451-455) purges queues/by_id alongside AdapterRegistry.free_slot, matching the earlier ack()-only cleanup path.

Beyond that, I spot-checked the load_state path and /adapter_runs yaml_path handling flagged as candidate issues this run: both do pass client-supplied paths to open()/restore without containment checks (unlike save_state's regex-validated tag), so that observation holds, but this control-plane HTTP server is driven by trusted internal tooling rather than arbitrary end users, which is a plausible reason it wasn't escalated. Given the size of this refactor (100+ files, new operation-contract/ledger/residency subsystems, checkpoint restore, HTTP surface) and its security-adjacent multi-tenant slot semantics, a human look remains worthwhile alongside this automated pass.

Extended reasoning...

This run's bug hunter reported no findings. This bot's own prior review on an earlier revision flagged two issues: an OperationLedger memory leak (fence/free_slot never purged queues/by_id, only ack() did) and a Dockerfile revert from a pinned Megatron-Bridge SHA to a mutable branch. I independently verified in the current diff that the ledger leak is fixed — fence() now pops queues/by_id (operations.py lines 451/455), matching commit f01778a ("bound the operation ledger by the registry's completed ring"). The Dockerfile still tracks @ bridge (not repinned to a SHA) despite the author resolving that thread without a code change, but that finding was only a nit-level build-reproducibility concern to begin with, not a functional bug, so restating it would be low value and risks being a broken record on a dismissed nit.

This run's candidate-issue list surfaced seven new items (rid/serving_name consistency, load_state/yaml_path path containment, checkpoint registration_id validation, SGLang KV-cache namespace collision, unused ttl_seconds) that were investigated and ruled out rather than reported as bugs. I spot-checked two of the path-handling ones directly in config.py/backend.py/http_server.py and confirmed the code shape described (load_state accepts any non-empty path string; yaml_path goes straight to open()/yaml.safe_load via parse_adapter_run_yaml) is accurate — no containment logic exists for either. I did not find independent evidence of exploitability given the apparent trusted-operator threat model of this control-plane service, consistent with the sub-agents' decision not to escalate these as bugs.

Given the PR's scope (a 100+ file refactor introducing new operation-contract, ledger, residency, and executor subsystems plus a service HTTP surface with client-supplied paths and checkpoint restore across multi-tenant LoRA slots) and that this is not a mechanical or narrow change, the appropriate level of scrutiny remains high, warranting continued human review even though no new bug was found this run and the previously-identified real defect has been fixed.