Skip to content

Add precision-aware model loading and MiniMax H3 training - #219

Open
Jayce-Ping wants to merge 26 commits into
mainfrom
feat/precision-aware-model-loading
Open

Add precision-aware model loading and MiniMax H3 training#219
Jayce-Ping wants to merge 26 commits into
mainfrom
feat/precision-aware-model-loading

Conversation

@Jayce-Ping

@Jayce-Ping Jayce-Ping commented Aug 28, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • add released Diffusers MiniMax H3 T2VA/FL2VA/Ref2VA workflows with structured video/audio replay and precision-safe timestep handling
  • separate component load dtype from trainable/frozen storage dtype, preserving model-declared FP32 islands
  • plan loading by physical component roots and coordinate target, auxiliary, reward, and host resources through backend-owned lifecycles
  • keep prepared target routes backend-managed while moving only frozen sibling state in composite roots such as Bagel
  • add FSDP2 selective-loading capability gates, cached replica fingerprint checks, and compatibility fixes for Bagel, SenseNova, Wan, SD3.5, and Z-Image
  • enable Bagel TDM through the shared batched media-free decode contract, adapter-owned reference guidance mapping, and integral scheduler-coordinate support
  • add a GPU-validated Bagel TDM recipe with ODE replay, stable NaViT batch composition, and generator/fake role optimizers

Test plan

  • python -m pytest -q — 1722 passed, 14 skipped
  • 32-GPU MiniMax-H3 T2VA + GRPO + FSDP2 full fine-tuning, 1 epoch with finite rewards/losses/gradients
  • MiniMax-H3 LoRA on DDP, DeepSpeed ZeRO-2, and FSDP2
  • SD3.5 LoRA on DDP, DeepSpeed ZeRO-2, and FSDP2
  • Z-Image LoRA on DDP/ZeRO-2/FSDP2 and full fine-tuning on FSDP2
  • real-weight loading/dtype inventories for Bagel, SenseNova, Wan2.1, MiniMax-H3, SD3.5, and Z-Image
  • focused loading contracts for LTX2, Wan2.2, and Qwen-Image-Edit-Plus
  • 32-GPU Bagel multi-reference I2I + GRPO + DeepSpeed ZeRO-2 reached W&B logging, initial eval, and training with finite PickScore rewards
  • 32-GPU Bagel TDM B=2 real-weight smoke completed one epoch with finite fake/generator losses and gradients
  • 64-GPU Bagel TDM long run launched with 256 unique samples/epoch, B=2, and group size 1; initial eval and multiple training epochs are finite (W&B)

Copilot AI lite review requested due to automatic review settings August 28, 2026 06:45

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

Let DMD2, TDM, and TDM-R1 preserve adapter-specific inference parameters through extra_kwargs while continuing to reject known retired algorithm controls.
Make video frame geometry explicit across training algorithms, remove DMD2's custom argument rejection path, and require the released diffusers 0.40 MiniMax H3 API.
Construct MiniMaxH3Blocks explicitly so diffusers 0.40.0 can load local modular checkpoints without its unresolved default-blocks path.
Bypass diffusers 0.40.0's broken local from_pretrained blocks forwarding while retaining its public constructor and component-spec loading behavior.
Use the official workflow-pruned from_pretrained path unchanged for both Hugging Face repo IDs and local unified checkpoint directories, removing redundant spec mutation and partition validation.
Keep the official ModularPipeline loader for both Hub IDs and local directories, rebinding only local ComponentSpec sources so component loading stays offline.
Convert the B=1 empty decoded-video batch to a per-sample None before constructing structured H3 samples for TDM and DMD2 rollouts.
Avoid overriding Diffusers-managed FP32 islands in the validated T2VA recipes.
Separate component load dtype from storage precision and centralize target, auxiliary, and reward loading so distributed backends cannot mis-handle aliases or replicated weights.
Prevent target-owned roots from re-entering auxiliary device management, make FSDP validation collective-safe, and remove unused loading-plan metadata.
Keep prepared logical targets in place while moving sibling modules needed by composite pseudo-pipelines such as Bagel.
Align runtime documentation with Bagel's partial-root residency behavior and guard fully excluded target roots from movement.
Point both SenseNova training examples at the intended SFT initialization.
Enable velocity-based KL regularization and use 14 denoising steps consistently in SenseNova examples.
Keep only the SFT model path update while deferring experimental KL and sampling settings until validated.
Use the shared vid_prompt dataset and align general GRPO, reward, EMA, and optimizer settings with the LTX2 T2AV recipe while preserving H3-specific geometry and guidance constraints.
Keep preprocessing at B=1 because MiniMax H3 does not accept batched outer workflow inputs.
Reuse the batched media-free decode contract for Bagel, map reference guidance through adapters, and accept integral scheduler coordinates so TDM can replay native trajectories without model-specific patches.
Flatten trajectory boundaries into ordinary backend accumulation work items so each graph is released after backward while preserving rollout-level feedback and optimizer-step semantics.
Allow backward-compatible full or selective block policies, keep checkpoint ownership singular under FSDP, and expose an ordered MiniMax H3 checkpoint-unit contract for memory/performance tuning.
Decouple evaluation sampling and media logging from optional reward computation so reward-free distillation runs still publish generated outputs.
Publish the working 512x768x124 configuration with sanitized save paths, noting that training needs over 200 steps to approach convergence.
@Jayce-Ping
Jayce-Ping force-pushed the feat/precision-aware-model-loading branch from 2d5c477 to 9c8331d Compare August 31, 2026 00:42
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.

2 participants