[models,trainers,hparams] refactor: infrastructure layer for multi-component models and multi-variant algorithms - #209
Merged
Conversation
Prevent duplicate state arguments and support partial active-count overrides while preserving legacy reduction semantics.
Treat an empty active-count mapping as no override while documenting forward-state key ownership.
…ajectories Migrate GRPO, GRPO-Guard, and DPPO off manual index-map arithmetic and onto the structured trajectory hooks, so heterogeneous components can be replayed with their own schedulers while single-latent runs keep their exact legacy numbers. Foundation: - Add IndexedTrajectoryTensor for compact per-step tensors behind sparse -1 index maps, with batched stack and device traversal. - Let StructuredTrajectory carry per-component log probabilities and named per-component callback trajectories in authoritative component order. - Expose per-component std_dev_t, dt, and log probabilities on ReplayStep and MultiModalStepOutput; the legacy bridge wraps each statistic under "latent" without recomputing or rescaling it. - Add BaseAdapter.get_replay_callback, get_state_active_numel, and get_train_step_indices, the latter rejecting misaligned scheduler groups. - Fix the legacy replay bridge to fall back to t_next = 0 on the terminal transition, matching what the trainers did before the migration. Trainers: - Replay through get_replay_step/get_replay_callback/forward_state and seed via set_trajectory_seed instead of reaching into the primary scheduler. - Build KL, Guard, and trust-region terms per component, then reduce them by active stochastic degrees of freedom rather than an unweighted component mean; a lone component is returned untouched so legacy values are bit-identical. - DPPO resolves effective sigma from each component's own dynamics_type. - Reference forwards keep KL arithmetic outside no_grad and each forward inside its own autocast region.
…omponent elements Scheduler statistics arrive as (B, 1, ...) broadcast tensors, so the GRPO-Guard scalar ratio silently formed a batch cross-product instead of a per-sample scale. Normalize each statistic to (B,) with component/field/shape context on failure. Global reductions (reference KL, DPPO trust-region KL) now consume raw per-element errors so every element is weighted once; a pre-reduced component mean rescaled by an active count cannot recover a masked sum. Guard keeps its per-component mse through a new overridable `reduce_component_latent_values` adapter hook that masked adapters can restrict to active positions. Also validate the current-policy joint log probability before PPO arithmetic in GRPO/Guard/DPPO, tighten trajectory dataclass and stack invariants with sample, field and component context, and build `return_fields` in a canonical order.
…and full replay invariants Latent reduction becomes a public validated wrapper over a protected overridable hook, for both the global and the per-component path. Both wrappers validate the input contract and the hook result, so an adapter override cannot bypass validation: the per-component result must use the authoritative component order, one (B,) tensor per component, and a shared batch size, device and dtype. Global reduction gains optional `state` context, passed from GRPO reference KL, DPPO trust-region KL and the Guard ratio, so a dynamic mask adapter can weight only the elements a per-sample state marks active. The default hook ignores state and stays bitwise unchanged. `_require_replay_log_prob` now requires a (B,) tensor sized by `replay.state` and reports trainer, step index, expected batch and received type/shape. `ReplayStep` validates every state component, every `ComponentTimes` mapping (allowing the legacy 0-dim terminal coordinate), the joint log probability and the component log probabilities. `MultiModalStepOutput` derives its authoritative contract from the first available mapping when no latent state is present, and requires `std_dev_t`/`dt` to be scalar-like per sample.
…ries DiffusionNFT, AWM and DPO now read their terminal state, component times, forward-process noise, and forward velocity through the trajectory hooks instead of indexing `all_latents[:, -1]` and calling `adapter.forward` directly. Every per-element reduction goes through the validated, state-aware reduction wrappers, so a component is weighted by its own active degrees of freedom and a dynamic mask adapter sees the per-sample state. Forward-process noising splits into an RNG-owning draw and an RNG-free application. `add_forward_process_noise` draws once with diffusers `randn_tensor` in `trajectory_component_order` and delegates to the new `apply_forward_process_noise`, which interpolates each component with its own sigma. That split is what lets DPO noise both preference arms from one shared noise tensor and lets the NFT/AWM precompute pass hand its exact noised state to the policy, reference and EMA passes without a second draw. The new `build_training_component_times` maps the single sampled scheduler coordinate onto per-component timesteps and sigmas without consuming randomness, so each trainer keeps its existing time distribution. Single-component behavior is bit-identical: NFT's normalized x0 matching and velocity KL, all five AWM weighting schemes and its PPO ratio, and DPO's velocity MSE, preference loss and implicit reward metrics are asserted against reproductions of the pre-migration formulas.
…nd validate velocity states The forward_state bridge now owns `advantage` alongside trajectory storage, so adapters that accept **kwargs no longer receive loss metadata that legacy decoupled trainers deliberately excluded, and an explicit bridge-owned kwarg fails fast instead of shadowing the batch. `require_velocity_state` now validates every velocity component against the state it was predicted for (exact shape, device, floating point, one shared dtype), so a broadcastable (B, 1, ...) velocity can no longer silently rescale the NFT/AWM/DPO matching losses. Latent states are checked for a leading batch dimension before their batch size is read, in the noising bridge, the shared forward helper and the DPO preference pair.
…ship wrapper `BaseAdapter.forward_state` now only owns the shared boundary — collision checks for state-owned and bridge-owned kwargs, plus batch filtering of trajectory storage and trainer metadata — and dispatches to a protected `_forward_state` hook for component-specific packing, forward and unpacking. The default hook keeps the single-`latent` numerics unchanged, so heterogeneous adapters override the hook and receive kwargs that are already clean. `__init_subclass__` rejects an adapter that overrides the public wrapper at class creation, naming the protected hook, so the ownership boundary cannot be bypassed accidentally.
Move DGPO and CRD onto the terminal-state, component-time, explicit noising, forward-state and state-aware reduction contracts. DGPO now resolves its clean state through get_terminal_state, builds component times from the shared timestep, and rebuilds its per-unique-id shared noise as a LatentState drawn in trajectory_component_order. The primary component keeps the legacy seed tuple and randn_tensor call, so a single-latent adapter stays bit-identical; further components extend the namespace by their declared index. DSM, reference DSM, clipping and KL all reduce through reduce_latent_values with the noised state. CRD keeps its two-pass lifecycle but pass 1 now stores only component times, the drawn noise and the detached old velocity; pass 2 replays them through apply_forward_process_noise and consumes no new RNG. The adaptive implicit reward normalizes each component with the validated component reducer before the global state-aware reduction, preserving the legacy elementwise-divide-before-mean order. Both trainers dispatch their epoch seed through set_trajectory_seed so every component scheduler is seeded, and both forward through forward_velocity_state instead of hand-built forward kwargs.
The tests written for the DGPO/CRD trajectory migration were left out of e509fbd because `tests/` is git-ignored; force-add them so the migration ships with the suite that gates it, matching the already-tracked Task 1-3B test files. Covers DGPO legacy shared-noise generator and tensor parity, unique-id sharing, declared component draw order, noising hook selection and draw counts, DSM/group/clipping/KL parity plus active-DOF reduction, and CRD zero-extra-draw pass-2 replay, adaptive and non-adaptive implicit reward parity, per-component normalization, every weight_temp branch and KL scaling.
…ifted flags Address the Task 3C review: the migration was pinned helper by helper, but nothing drove the real optimize() entry points, so gradient ownership, parameter-swap order, optimizer cadence and the distributed group reduction were only covered indirectly. Add a production-path suite that runs DGPOTrainer.optimize() and CRDTrainer.optimize() end to end against a tiny trainable adapter with per-scope biases, recording accelerator, optimizer and logger. It asserts that only the current-policy forward owns the graph, the mode/scope/backward/ clip/step/log order, the logged metric keys, the ema_ref and CRD snapshot decay updates, and CRD's pass-1-before-pass-2 lifecycle. Two independent legacy oracles reproduce the pre-migration one-batch/one-timestep formulas from raw tensors and compare loss, trainable-parameter gradient, RNG state and draw cadence with torch.equal. A peer-summing accelerator drives the real group-sum collective and checks the sigmoid consumes the globally summed contributions, and row-order permutations prove shared noise depends only on (unique_id, component). Production changes the tests forced out: - Replace the two `assert` invariants on the DGPO reference velocity with an extracted `_select_dgpo_reference` that raises with the gating flags (use_ema_ref, _requires_ema_ref, clip_dsm, clip_kl, kl_beta) instead of a bare truth check. - Thread inner_epoch, timestep index and the expected component order into every shared-noise validation error, and validate the assembled noise state order before returning it. - Validate CRD pass-1 noise against the reloaded terminal state before reapplying it, naming the timestep the drift belongs to. - Document why `component_name` stays in the draw hook's keyword contract even though the namespace keys off the authoritative component index.
…ates Prove with two micro-batches and two timesteps that every pass-1 snapshot forward completes before the train switch and before any policy/reference forward, and pin the full backward/clip/step/zero-grad/log cadence across them. Reject a pass-2 clean state or pass-1 stored noise whose component order drifts with a CRD+timestep expected/received message instead of a bare KeyError from the component lookup.
Replay both DiffusionOPD passes through get_replay_step/forward_state and structured component targets instead of legacy index maps and direct adapter.forward. - add resolve_scheduler_group_dynamics: validate every component scheduler, reject a mixed ODE/SDE group, and derive _is_sde from the group - add project_distillation_target_state and an optional explicit sigma on the single-tensor helper; x0 uses each component's stored sigma and keeps the flow-matching fallback for the legacy single-latent replay - add compute_structured_distillation_loss: one global detached self-normalization scale, per-component KL denominators applied inside their own component, and an active-DOF weighted global reduction with state context - store the teacher cache as an ordered component mapping and validate its order, batch size and stored step count before PASS 2 indexes it - derive each SDE denominator from its own component scheduler and request std_dev_t/dt only for the stochastic student pass - seed every component scheduler via set_trajectory_seed in start() Single-component ODE and SDE losses and gradients stay bit-identical to the pre-migration formulas.
…educe The stochastic structured loss reduced each component to a per-sample mean, divided that mean by the component KL denominator, and then called the adapter reducer a second time. A masked adapter reducer expects latent-shaped values, so the second call fed it pre-reduced scalars and silently produced a loss that no longer matched the active elements. Denominators now divide the raw squared errors inside their own component and exactly one reduce_latent_values runs over the scaled elements; global self-normalization stays a per-sample scalar on that single reduction, which keeps the legacy arithmetic bit-identical. Also validate the whole ComponentTimes contract (authoritative component order, one value per sample, state device, and the single documented legacy terminal scalar-zero t_next) before any target branch, reject a rank-0 KL denominator instead of indexing it, and prefix projection errors with the teacher/student pass and replay step so a two-pass failure is attributable.
Dividing raw elements by the KL denominator is only bit-identical to the legacy formula when the denominator is a power of two, and real schedulers produce arbitrary positive transition variances. A single component now divides the reduced (B,) value after self-normalization, exactly as the legacy trainer did, which stays correct under a dynamic mask because the denominator is a per-sample scalar. Several components may carry different denominators, so those still scale the raw squared errors of their own component before the one shared state-aware reduction. Also validate that every structured latent component is batched and shares one batch size before reading shape[0], and re-check denominator finiteness and positivity in the trainer so an unusable transition variance names its replay step and component.
Add an optional static active-mask contract to ComponentTrajectory and
LatentState so fixed conditioning tokens can be stored without being treated
as stochastic degrees of freedom, and honour it in the trajectory bridge
(replay/terminal/callback, forward-process noising, active counts, reducers).
Legacy unmasked states keep their previous behaviour.
Make LTX2 T2AV and I2AV the first adapters with a two-component
("video", "audio") order: a twin-scheduler SchedulerGroup built during adapter
init, mirrored training times, ordered video-then-audio noise draws, and a
protected component forward hook that exposes real per-modality latents,
statistics and log-probs. The legacy concatenated forward() keeps its API and
numerical outputs.
…eline coverage Validate every component input the joint LTX2 forward consumes before it concatenates or calls the model: per-sample time coordinates on the state device, matching batch/channel/dtype/device across video and audio, stored next-state components matching their current counterpart exactly, and stored next-state active masks matching the current static masks. Errors name the adapter, field, component and the expected/received values. Add the test coverage the review found missing: the real LTX2 constructor path (with only pipeline loading, target-module parsing, freezing and precision casting patched) proving BaseAdapter.__init__ builds both twins in video-then- audio order with no post-super duplicate, a baseline oracle pinning the public concatenated forward against outputs captured from the pre-Task-4A implementation at ee6d247, and lifecycle dispatch order across eval/train/ rollout/set_seed.
…venance I2AV now runs the shared component/time contract before reading the video active mask, so a missing or reordered component reports the contextual expected/received error instead of a KeyError. The legacy-forward golden fixture is regenerated by a tracked script that refuses any checkout whose HEAD is not the oracle commit, and the JSON now records the full commit plus both adapter blob hashes, which the tests re-verify against Git objects. The saved RNG state sum is compared exactly.
The golden identity is now pinned to fixed full commit/blob constants and every output, dispatch-order and RNG comparison reads only the JSON, so all 12 cases run from a source export or a shallow clone. The live rev-parse cross-check resolves objects through a lookup that reports unavailability instead of failing, and skips when Git, the repository or the referenced objects are missing. The tracked generator stays strict and Git-dependent.
Rewrite the LTX2 T2AV/I2AV inference collectors so a rollout publishes one authoritative StructuredTrajectory per sample (per-component states, full per-component schedules, joint and per-component log probabilities, and the latent-shaped callbacks) and leaves every legacy trajectory field None. I2AV carries a video active mask derived from the conditioning mask, so the conditioning frame stays out of every reduction, log-prob weighting and forward-process noising. Add the shared builder to models/ltx2/_common.py, inference-capable fakes, an independent legacy-loop oracle, and tests covering builder contracts, full rollout/RNG/decode parity, stack/replay plus condition-frame invariants, and the algorithm interface matrix on a real LTX2 adapter instance.
Pin explicit-generator and non-default decode behavior against the legacy oracle so collector changes cannot alter either RNG stream or dispatch order.
Read declared components from ModularPipeline's public spec APIs so workflow pruning and named materialization remain correct before lazy values enter pipeline.components.
Use pretrained and config component-name APIs for complete lazy declaration discovery while keeping component_names scoped to currently registered values.
…e adapter Trainers reconstructed x0 with a trainer-local formula that assumed one velocity direction, so an adapter declaring the opposite convention silently got the wrong target. The bridge now owns the projection and the coupled and decoupled trainers consume the adapter's declared direction. Cherry-picked from the multimodal branch without its model-specific parts.
.scratch/ is git-ignored working space for temporary analysis; this LTX2 trajectory report was force-added into a commit and does not belong in the repository. The file stays on disk, it is just no longer tracked.
Jayce-Ping
force-pushed
the
feature/structured-trajectory-runtime
branch
from
August 18, 2026 10:27
1bf7a9e to
87755c7
Compare
…onfigures
Every shipped example that carries an `optimizers:` block names its single entry
`default`, while the trainable role is `base`. `_optimizer_args_for_role` fell back
to that entry unrenamed, and `build_optimizer` then looked its parameters up by
`OptimizerArguments.name` in a table keyed by role. The lookup missed, so all 54 of
those configs died at startup with
ValueError: expected optimizer 'default' to own parameters, received none
reproduced end to end on 8 GPUs under both DeepSpeed ZeRO-2 and FSDP, so it is not
backend-specific. Config parsing never caught it: the files load fine and the
mismatch only appears once roles and optimizers are joined.
The fallback now returns the entry renamed to the role it is configuring, which is
what "a single-policy run need not name it" already promised. A config that does
name its roles is untouched, and asking for a role no entry names still raises.
…dient The backend gate refused only DeepSpeed and told the reader to "use DDP or FSDP". FSDP1 does not work: it flattens each wrapped unit into one 1D FlatParameter, so `split_muon_parameters` hands Muon the matrices it sees before `prepare`, and the first optimizer step then fails with "Param gradient must be a 2D matrix" -- after a full rollout has already been generated. Measured on 8 GPUs with SD3.5 LoRA GRPO: FSDP1 (fsdp_grad_op_shard.yaml) crashed at the first step, FSDP2 (fsdp2.yaml) completed with ratio_min = ratio_max = 1 and ratio_std = 0. FSDP2 shards each parameter as a 2D DTensor, so rank survives `prepare` there and the composite still sees matrices. The gate now splits the two cases and reads `fsdp_plugin.fsdp_version`, so an FSDP1 Muon run is refused at startup with the config change that fixes it. FSDP1 stays available to AdamW, which accepts a flattened parameter.
…train
An 8-GPU SD3.5 LoRA run under FSDP1 exits 0, logs a plausible loss and changes no
weights. `prepare` reports
Total: 282,743,448 Trainable: 0 Percentage: 0.00%
against 18,776,064 trainable on DeepSpeed, DDP and FSDP2. FSDP1's transformer-based
auto-wrap flattens each block's frozen weights together with its adapter into one
FlatParameter, and a flattened unit carries a single requires_grad, which the frozen
majority wins. Everything downstream then behaves consistently with having nothing to
train: the gradient norm is exactly 0 while other backends report 6e-4, and saving the
LoRA checkpoint asserts inside FSDP because no adapter parameter is gathered.
Role ownership is already validated, but only before `prepare`, so it cannot see a
backend that absorbs the parameters while wrapping. The new check runs on the prepared
root and fails startup, naming FSDP2 as the FSDP that keeps LoRA trainable. It is
written backend-agnostically: any future wrapper that swallows the gradients stops the
run instead of burning a training budget on a no-op.
…e trainability across ranks Two mistakes with one root: reading a sharded plan as if a rank saw the whole model. `clip_grad_norm_` received `get_trainable_parameters()`. Accelerate only delegates to FSDP's collective `clip_grad_norm_` when the list it is handed is exactly `model.parameters()`; given a subset it falls through to the plain utility, which under FSDP1 computes the norm from this rank's shard alone. Every rank then clipped by a different norm, and the reported value on rank 0 was 0 because the LoRA slice lives on another rank. Passing the prepared root's parameters restores the collective path; frozen parameters carry no gradient, so nothing changes for DDP or DeepSpeed. Measured on 8 GPUs with SD3.5 LoRA: FSDP1 grad_norm 0 -> 6e-4 for GRPO and 0 -> 7.26e-2 for NFT, both now identical to ZeRO-2, DDP and FSDP2. The guard added in bbdfe48 was rank-local for the same reason and was wrong. FSDP shards a flattened unit by byte range, so with LoRA the frozen base fills the early shards and the whole adapter can land on one rank -- an 8-rank repro puts all 2048 trainable elements on rank 7 and none on ranks 0-6. Counting locally rejected runs that train correctly. The check now reduces across ranks and fires only when no rank has anything to train, which is the invariant that was meant.
Saving a LoRA checkpoint under FSDP1 poisoned the next rollout. The save called `FSDP.state_dict_type(component, ...)`, but components live inside the prepared bundle, so `component` is a non-root FSDP instance; entering that context lazy-initialized it as a root, and the following forward died in the real root's `_root_pre_forward` with "Non-root FSDP instance's `_is_root` should not have been set yet" -- an epoch after the save that caused it. Switching to `get_model_state_dict` alone did not help: any state-dict API does the same lazy-init to whatever module it is given. The gather now goes through the root that `prepare` returned and slices the requested member back out by its `members.<name>.` prefix, which is how the bundle's ModuleDict names them. `FSDP.state_dict_type` also carried a deprecation warning per wrapped block; the modern API replaces it, matching what the FSDP2 branch already did. `ignore_frozen_params` is no longer forwarded. Torch removes frozen entries with `state_dict.pop(fqn)` and no default, and under FSDP1 the FQN it rebuilds for a PEFT parameter wrapped in both `_fsdp_wrapped_module` and `_checkpoint_wrapped_module` is absent from the gathered dict, raising KeyError on a frozen base weight. The existing `state_dict_keys` filter already narrows the result to the adapter. Verified on 8 GPUs: FSDP1 SD3.5 LoRA now writes checkpoints and its weight delta matches DeepSpeed's -- 190 of 382 adapter tensors changed, max 3.03e-04 against ZeRO-2's 2.99e-04 -- so the adapter is genuinely being trained and saved.
Collaborator
Author
Verification complete — ready to mergeEverything in the test plan above now has a result. Six defects surfaced only under real
One correction worth recording: Remaining known limitations, none blocking:
|
…y wrapper identity The multi-role backend check required `accelerator._models[0] is self.model_bundle`. Accelerate registers the module in `_models` before wrapping it, so under DDP the tracked entry is the inner bundle while `prepare` returns the DistributedDataParallel around it. Comparing wrappers therefore rejected every multi-role DDP run with "expected the tracked prepared model root to be self.model_bundle, received different identities", which is how an 8-GPU DMD2 run failed at startup. DeepSpeed passed only because its engine happens to be what gets registered. Comparing unwrapped identities states the invariant that actually matters -- the single tracked root is the one this trainer drives -- and holds for DDP, FSDP and DeepSpeed alike.
A checkpoint reconstructed its component layout in four independent places -- save_checkpoint, _load_lora, _load_full_model and _detect_checkpoint_type -- so nothing kept the writer and the readers agreeing. Route all four through one _checkpoint_entries resolver and record what it produced in a manifest, which becomes authoritative when present and is reconstructed from the legacy rule when it is not, so existing checkpoints keep loading. Multi-role algorithms were the ones paying for the drift. Only the base variant was ever written, so a DMD2 run that saved every five epochs and later resumed came back with its generator restored beside a freshly initialized fake score, trained against the wrong critic and reported nothing wrong. A periodic checkpoint now carries the training-only roles too, each in its own directory that PEFT can read directly, and each loads back into its own variant instead of whichever adapter happened to be active. Initializing from a base-only export stays legal and now says which roles it could not restore. Full-weight loads stopped replacing the component object: it is a member of the prepared root and a variant of the registry, and swapping it detached both.
…s exist The adapter restores weights while it is being built, and the trainer declares the roles it trains afterwards. A multi-role resume therefore arrived before its variants did: both roles were loaded, both routed to the single live adapter, and the fake score silently overwrote the generator it was supposed to sit beside. Place only the primary artifact during that early load, and let the trainer finish the restore right after it declares the variants.
_mix_precision runs while the adapter is built and only sees the base components. The extra variants are materialized afterwards, from PEFT's fp32 defaults, and nothing cast them: a DMD2 checkpoint shipped a bf16 generator beside an fp32 fake score. That is twice the memory, a different numerical path for the two halves of the objective, and an RMSNorm whose weight no longer matches its activations, so every norm the fake score feeds falls off the fused kernel and reports it once per rank as a dtype mismatch warning.
Collaborator
Author
Checkpoint layout, and the resume it was quietly breakingThree commits since the last update, all in the same area. A checkpoint reconstructed its directory layout in four independent places -- Multi-role algorithms were paying for the drift. Only the base variant was ever written, Layout, unchanged for the single-component single-role case that every released checkpoint Verification
Mergeable: no conflicts with |
8 tasks
Two roles can be swapped and still look plausible; three cannot. TDM-R1 ships a generator, a fake score and a surrogate, so cover that shape and assert no two roles collapsed onto the same weights.
… is asked for BaseAdapter._init_ref_parameters reads training_args.ref_param_device, and every algorithm inherits it, but only six declared the field. The rest died with an AttributeError the moment a full-weight run needed a reference -- which is exactly what the distillation trainers do. Move it to the base and drop the six copies.
…ckpoint Re-running an experiment under the same run_name writes into a directory that already holds another model's artifacts, and a save only overwrites what it happens to produce. A shard index is the leftover that hurts: the loader trusts it ahead of a single-file save, so an index written by a 9B model sent a 2B resume looking for shards nobody wrote. Clear an entry's own artifacts before writing it, leaving the roles and components nested inside for their own turn.
…context closes PEFT keeps only the active adapter trainable, and leaving disable_adapter re-marks only that one. A multi-role objective queries its frozen reference between a role's forward and its backward, so the role it was training came out of that context frozen: measured on TDM-R1, both non-base roles went from 382 trainable parameters to 0 across the query. Autograd reads requires_grad when backward executes rather than when the graph was built, so every gradient was dropped and the optimizer stepped on nothing, with no error anywhere. Reassert the invariant after the adapter contexts unwind -- after, because leaving one is itself what narrows it.
Distillation may skip expensive media reconstruction while retaining trajectories and conditioning metadata. Expose a model-agnostic adapter hook for the empty decode result so trainers never inspect model type, modalities, or conditioning fields; heterogeneous adapters can include additional decode metadata in their override.
Let adapters opt into an ordered image/video/audio manifest without exposing conditioning semantics to trainers. Canonicalize and hash source manifests, decode references with row/index-aware failures and preserved FPS/sample-rate overrides, pass transient media only into preprocessing, and cache only Arrow-safe outputs plus the canonical identity manifest. PyAV remains an actionable optional boundary until a reference-capable model is installed.
Target components are required lifecycle participants, not optional lazy declarations. Resolve each through ComponentRuntime when enabling checkpointing so an invalid configured name fails with declared-component context instead of being silently skipped.
Include the resolved dataset root, source-content digest, adapter-declared geometry fields, and preprocessing cache version in cache fingerprints so reference workflows cannot reuse stale embeddings across manifests or layouts. Let generic audio loading fall through when modern torchaudio lacks its optional torchcodec decoder.
8 tasks
Move replay, state, batching, backend, and checkpoint contracts out of algorithm trainers so their orchestration stays focused without changing the public trainer hooks.
Preserve heterogeneous metric dtypes bit-for-bit while collapsing eval and CRD gathers, and carry zero-variance counts in the existing statistics reduction.
Pack same-dtype sample fields when A/B shows a win, while retaining field-wise gathers for mixed dtypes and large CPU-bound payloads that regress latency.
Document why individual trainable roles cannot be manually offloaded after prepare and direct users to safe snapshot, component, sample, and backend-managed controls.
Let one frozen dtype policy preserve checkpoint precision by default while overriding component groups or concrete modules, without breaking scalar configs or FSDP2 master dtype invariants.
Normalize scheduler coordinates through float64 so the largest float32 timestep below 1000 cannot round back to sigma one on CUDA and invalidate TDM conditional re-noising.
Sanitize portable defaults and keep optimizer, trainer, dtype, and backend documentation consistent with the refactored runtime.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Infrastructure layer for multi-component models and multi-variant algorithms. Model code owns caller-named component variants and snapshots; trainers own role meaning, scheduling, optimization, and checkpoint compatibility.
What lands
ModelBundle, and role-aware checkpoint manifests.BaseTrainer.gather_samplespacking.frozen_parameters_dtypeaccepts a scalar or default/group/component mapping, allowing checkpoint-native VAE precision without changing transformer or text-encoder policy.Compatibility and performance
Verification
transformers/text_encoders=bf16,vae=fp32, and exact on-policy replay.Stack
Merge order: #209 → #210 → #212.
Current head:
98b6b4d.