Skip to content

feat: add on-demand models with shared LoRA and exclusive training - #312

Draft
Yunnglin wants to merge 2 commits into
mainfrom
feat/on-demand-training
Draft

Yunnglin wants to merge 2 commits into
mainfrom
feat/on-demand-training

Conversation

@Yunnglin

@Yunnglin Yunnglin commented Oct 13, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

Expose a catalog of trainable models without keeping all model workers resident.
A training-handle request allocates a single-node Ray placement group and starts
an instance on demand. LoRA tenants share an instance; full-model tenants own
exclusive instances. Business-idle handles expire, and a dynamic instance releases
its devices after its final tenant leaves, allowing another model to use them.
Configured resident instances remain available and reset tenant state before reuse.

Keep the existing applications YAML entry point and one training controller.
idle_timeout_s measures business inactivity; session-heartbeat expiry and maximum
lifetime retain their existing separate meanings. Capacity exhaustion is rejected
without adding a waiting queue or preemption. Recreating an expired handle starts
clean state; the lifecycle does not implicitly save or resume a checkpoint.

Standalone scope

This PR is based directly on main and can merge without #315. The backend
changes are limited to resetting resident full-training state/hooks and using the
supported PretrainedConfig import for Megatron startup. Checkpoint restoration,
optimizer master-copy/rebinding repairs, tied-weight loading, and Megatron TP/SP
communication fixes remain in #315. The corresponding implementations here match
main, including the original lazy optimizer wrapping.

The controller preserves tenant isolation, startup/cleanup ownership, epoch fencing,
quota recovery, bounded retry receipts, and operation tracking through the compute
queue. A timed-out runtime stays quarantined even if its actors answer ping. Queue
integration preserves the sampling concurrency and cancellation behavior already on
main. The example configuration and committed device acceptance now use ordinary DP.

Local validation

  • Server/model/Transformers/client CPU regression: 828 passed, 0 failures/errors;
    96 hardware or external-service/gated cases skipped in that broad run. The
    applicable Redis and YAML gates were exercised separately below.
  • Redis/state/receipt/recovery regression: 30 passed, 0 failures/errors/skips.
  • Focused training queue, actor-health, and existing sampling-concurrency regression:
    64 passed, 0 failures/errors/skips (overlaps the broad suite).
  • Installed Native/Tinker SDK HTTP mock acceptance passed: shared LoRA, exclusive
    full training, release/recreation, runtime-route rejection, controller restart,
    state loss, and ten reuse cycles. Real Ray allocation and fragmentation probes
    also passed with logical CPU resources; these are not device-training evidence.
  • Existing applications YAML startup passed with both memory and Redis state.
    The mock checkpoint fixture supplies its LoRA configuration for the checkpoint
    classification introduced on main; production sampler validation is unchanged.
  • Repository pre-commit checks passed.

Standalone DP acceptance on PPU

Source: dd96d85464b63e37d195820393785bbf98cf93b6, directly based on main 444bd8e82534e484f7e5296b390409727a68bb0e.
The #315 commits are not ancestors of this head. Validation ran the committed
tools/validation/on_demand_training/device_validation.py in a clean isolated
checkout, with verified source hashes, import provenance, environment, and collection.

Backend Strategy Result
Transformers Accelerate ordinary DDP 4 passed, 0 failures/errors/skips
Megatron MCore DDP 4 passed, 0 failures/errors/skips

Both backends use DP=2, TP=1, no FSDP, two devices per instance and four visible
devices per isolated Ray/Serve cluster. Every device holds complete model weights.
Transformers uses Accelerate ordinary DDP. Megatron uses MCore DDP and retains its
default full-model distributed optimizer, which shards optimizer state, not model
weights; shared LoRA does not use the distributed optimizer. Models are cached
Qwen2.5-0.5B-Instruct and Qwen3-1.7B.

Every Native/Tinker × LoRA/full case asserts lazy allocation, actual updates,
tenant isolation, both users training, continued training after a peer releases,
release of devices, cross-model reuse, idle reclamation despite heartbeats,
expired-handle rejection, and clean recreation. LoRA cases also assert expansion
when slots fill and rejection when the device pool is exhausted. Deep health,
training logprobs, placement/device identities, exit status, JUnit, and cleanup are
checked. These correctness fixtures do not measure model quality or cold-start SLA.

Both detached jobs exited with status 0. Task-owned processes were cleaned up,
all eight devices returned to the 2 MiB baseline with no device processes, the
pre-existing service remained accessible, and the selected checkout stayed clean.

Limits

This DP suite performs no checkpoint I/O. The pre-existing full
Qwen/Accelerate FSDP2 Adam-export KeyError: 0 is unfixed and outside this resource
lifecycle change. Checkpoint/FSDP/TP/SP results from the previous stacked revision
are not acceptance evidence for this standalone head. Resident reset and fault
quarantine have CPU/mock coverage; they were not rerun on devices in this DP suite.
No new FSDP/TP/SP, MoE/FP8/Hybrid, or multi-node acceptance is claimed.

Raw environment/configuration, source and script hashes, collected cases, logs,
JUnit, numerical results and cleanup evidence are retained outside the public
source tree. Shared PPU packages were not changed. Manual PPU evidence is separate
from GitHub CI. Keep this PR in draft.

@Yunnglin
Yunnglin marked this pull request as draft October 14, 2026 00:40
@Yunnglin Yunnglin changed the title feat: 支持小模型按需加载、共享 LoRA 与独占训练 feat: add on-demand models with shared LoRA and exclusive training Oct 14, 2026
@Yunnglin
Yunnglin force-pushed the feat/on-demand-training branch from 7f1cbc1 to 4d8888b Compare October 14, 2026 06:38
@Yunnglin
Yunnglin changed the base branch from main to fix/training-state-restoration October 14, 2026 06:44
@Yunnglin
Yunnglin force-pushed the feat/on-demand-training branch from 4d8888b to 87fd772 Compare October 14, 2026 06:57
Preserve legacy expiration semantics, isolate user initialization failures, and keep recovery nonblocking. Bound retry receipts and reconcile lost completion notifications.
@Yunnglin
Yunnglin force-pushed the feat/on-demand-training branch from 87fd772 to dd96d85 Compare October 15, 2026 02:43
@Yunnglin
Yunnglin changed the base branch from fix/training-state-restoration to main October 15, 2026 02:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant