docs: complete data flow performance model for DreamZero training#81
Draft
zhangqin200182 wants to merge 3 commits into
Draft
docs: complete data flow performance model for DreamZero training#81zhangqin200182 wants to merge 3 commits into
zhangqin200182 wants to merge 3 commits into
Conversation
Systematic trace of data volume, shapes, and formats through the entire training pipeline: DROID dataset → DreamTransform grid composition → VAE/CLIP/T5 encoding → patch embedding → token assembly → DiT 40-layer forward → loss computation. Key metrics per training step (B=1, 8× Ascend 910): - Grid video input: 11.2 MB (3 views composited to 176×320) - VAE latent: 0.62 MB (45× compression) - DiT token sequence: 3985 tokens × 5120 dim = 38.9 MB (bf16) - Per-layer data movement: ~735 MB (FSDP amortized per card) - HCCL total communication: ~78.9 GB per step per card - Peak HBM: ~19.5 GB (32% of 61.28 GB capacity) - Trainable params: ~86M (0.52% of 16.5B total) Includes FSDP all-gather/reduce-scatter patterns, gradient checkpointing analysis, optimizer state sizing, and optimization opportunities. Co-Authored-By: Claude <[email protected]>
Corrections from deep code tracing: - Per-block params: 403.8M (not ~350M) — includes k_img, v_img in cross-attn - LoRA: 19.2M params, 400 adapters (10/block — PEFT matches q,k,v,o in BOTH self_attn AND cross_attn by suffix) - Total trainable: 113.9M (LoRA 19.2M + encoders/decoders 94.7M) - Per-block bf16 size: 807.7 MB, per-card shard: ~101 MB - HCCL total: ~148 GB bidirectional per step per card - Gradient checkpointing: use_reentrant=False (not True) - Added FSDP cleanup mechanism section (3-phase NPU workaround) Co-Authored-By: Claude <[email protected]>
Complete comparison of LoRA vs full training across all dimensions: Memory (per card, 8 NPU): - LoRA: ~9.4 GB (15% HBM) → Full: ~34.6 GB (56% HBM) - Main increase: optimizer states (+24.75 GB for 16.5B params) HCCL Communication: - LoRA: ~117 GB bidirectional → Full: ~231 GB bidirectional (1.98x) - Reduce-scatter dominates: 57.7 GB vs 0.5 GB per card Step Time: - LoRA: ~32s → Full: ~40-55s (1.25-1.7x slower) - Optimizer step dominates: 5-8s (16.5B params) vs 0.5-1s (114M) Also covers: - 16 NPU scenario (33% HBM, lower per-card comm) - Checkpoint strategy (sharded: 8.25 GB/card, full: 66 GB) - Decision matrix: LoRA vs Full trade-offs - Data flow unchanged vs changed stages Co-Authored-By: Claude <[email protected]>
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
1 participant
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
Systematic trace of data volume, shapes, and formats through the entire DreamZero training pipeline, from DROID dataset to DiT forward pass.
Key Metrics (B=1, 8× Ascend 910 NPU)
Contents
🤖 Generated with Claude Code