Skip to content

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

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

Yunnglin wants to merge 11 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
- Split runtime allocation from deployment so the controller persists
  ownership before any worker starts; the driver no longer writes records.
- Make cleanup, close and failure transitions persist-before-mutate, track
  failed instances for sweep retries, fence admission, and exclude failed
  residents when rebuilding minimum replicas.
- Never lose a worker wake-up when enqueue callbacks or queued-status
  writes fail; log background job failures with tracebacks.
- Return 503/Server while the controller is not ready; drop unreleased
  reason codes; reset resident training state before model-id branching.
- Move generic device allocation to server/resources, managed runtime to
  server/model, training records/store to server/state, TaskQueueConfig to
  server/config, and shared proxy/env propagation to the server root, with
  no compatibility re-exports.
- Encapsulate ServerState backend access and batch heartbeat reads outside
  the controller lock.
- Guard all server component dependencies with a strict grimp layer test.
- Require ray[serve]>=2.55 for async deployment constructors.
- Stop tracking local-only docs/design notes.
Fix optimizer metrics, wrapped parameter groups and safe Redis reconnects. Add lifecycle metrics, isolated Redis/GPU acceptance and stress coverage. Move shared probes into tests and untrack local tools.
Integrate 015732e while preserving controller-owned training creation and sequence ledgers. Adapt full-mode regression tests to the catalog/controller path. Verified 434 targeted tests and the complete pre-commit suite.
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