Repository navigation
Speed up OpenWAM-α fine-tuning 2.4-3.1x with ResearchBoost-Agent - #58
Merged
xueminchi merged 15 commits intoNov 5, 2026
Merged
Conversation
With training.gradient_accumulation_steps (GAS) > 1 the accumulated gradient should be the mean over the micro-batches. accelerate leaves the 1/GAS scaling to DeepSpeed, but DeepSpeed 0.18.3 to 0.19.5 apply it only through backward hooks registered in DeepSpeedEngine.forward() (deepspeedai/DeepSpeed#7665, released in 0.18.3; #8245 in 0.19.6 moved the division back into engine.backward()). The trainer computes the loss through the architecture's own methods and never calls the engine's forward, so the gradient that is clipped and stepped is the sum over the micro-batches (4x the mean at GAS = 4). The pinned range deepspeed>=0.18.5,<0.19 is affected in full; GAS = 1, used by the shipped recipes, is not. Divide the loss by GAS before accelerator.backward() on exactly those DeepSpeed versions. GAS = 1, other DeepSpeed versions and runs without DeepSpeed get the loss unchanged, so nothing is scaled twice and GAS = 1 training stays bitwise identical. Logged losses are not scaled. LIBERO fine-tuning on 4 GPUs, 16 per GPU x GAS 4 against 64 per GPU x GAS 1 on the same samples: the grad-norm ratio was 3.94 before the fix and 0.99 with it.
The host work in front of the model (umT5 on every batch, per-frame float conversion and blocking host-to-device copies) is a large part of a fine-tuning step. New switches, all off by default: - OPENWAM_TEXT_CACHE=1: per-process cache of the frozen umT5 encoder's output rows, keyed by each sample's token ids and mask and filled inside training iterations (OPENWAM_TEXT_CACHE_MAX_ENTRIES rows of 4 MiB, default 256). LIBERO has about 40 prompts, so umT5 stops running after the first iterations. Cached rows are bitwise the encoder's. - OPENWAM_TEXT_MISS_FAST=1 (with the cache): on a miss, encode only the distinct missing prompts, on a prefix padded to a multiple of 64 tokens instead of 512. Same function; GEMM shapes differ, so rounding can. - OPENWAM_CONTEXT_TRUNC=1 (with the cache): drop the all-padding tail of the text context. Those positions are masked in every cross-attention. - OPENWAM_FAST_INPUT=1: collate frames into uint8 clip tensors in the DataLoader workers and convert them on the GPU in one batch with the same ops in the same order; copy small host tensors through pinned memory without blocking; memoize the RoPE table and the VAE latent statistics per device. The model inputs are bitwise the same. Applies to backbones that take uint8 clips (native Wan VAE with a variant that does not read frame pixels, such as TI2V); others keep list collation. - OPENWAM_VAE_FIRST=1: queue the VAE encode before the text encode, so tokenization and cache lookups overlap the VAE on the GPU. Tests compare the uint8 path with per-frame preprocessing and the cached text rows with direct encodes.
Outside the GEMMs, the MoT layers spend their time in masked attention, RoPE, sequence concatenation and chains of small elementwise kernels. New switches, all off by default: - OPENWAM_SDPA_CUDNN=1: masked scaled-dot-product attention asks for cuDNN's fused attention before PyTorch's default order. Masked calls otherwise run the memory-efficient kernel, about 2.5x slower forward + backward at these shapes. A different kernel, so rounding can differ. Deterministic mode keeps the default call: cuDNN attention has no deterministic implementation. - OPENWAM_FUSED_ROPE=1: Triton forward and backward for Wan's 3-D RoPE, with the float64 products and rounding steps of rope_apply, so the results are bitwise the same. The first call of a process compares with rope_apply and keeps rope_apply if they differ. - OPENWAM_FAST_CAT=1: the joint [video, action] q/k/v concatenation and the attention-output split as one Triton row-copy pass each, mirrored in the backward. Bitwise the same. - OPENWAM_COMPILE_BLOCKS=1: torch.compile the memory-bound chains of each Wan and action block half (AdaLN modulation, LayerNorm + modulate, q/k RMSNorm, gated residuals, tanh-GELU); GEMMs stay in cuBLAS. Same formulas with fused fp32 intermediates. Under deterministic algorithms the blocks use the eager modules, since the compiled chains are not reproducible from run to run there. GPU tests compare the Triton kernels with the PyTorch ops bitwise, forward and backward.
VideoVAE38_.encode runs the causal encoder chunk by chunk (the first frame, then 4-frame chunks) and carries cached frames between chunks. For a 1 + 4k frame clip that is the same function as running each layer once over the whole clip: the causal convolutions see the same windows, the stride-2 time convolution and the temporal average-pool see the same frame pairs, and everything else is per frame. New switches, all off by default: - OPENWAM_VAE_FULLCLIP=1: encode training clips in one pass. The convolutions run with other shapes, so the rounding can differ. - OPENWAM_COMPILE_VAE=1: torch.compile that encode, so Inductor fuses the norm, SiLU and padding chains between the convolutions. - OPENWAM_VAE_CL=1 (with the compiled encode): run the 3-D convolutions channels-last as opaque custom ops, which removes cuDNN's layout conversions around every convolution. The latents are returned in the default contiguous layout, so torch.randn_like(latents) draws the same noise for the same elements. A float64 test on a small randomly initialised VAE checks the whole-clip encode against the chunked encode.
Skipping activation recompute is one of the larger speedups, but with it 32 samples per GPU on 8 GPUs (global batch 256) needs a projected 120.4 GiB per GPU, more than an 80 GB GPU has. New switches, all off by default, trade memory against recompute or host copies without changing any computed value: - OPENWAM_CKPT_LAYERS=k: with use_gradient_checkpointing, checkpoint only the first k MoT layers (unset: all, as before). The other layers skip the recompute in the backward. - OPENWAM_MIXED_ATTN_CKPT=0/1: override mot_checkpoint_mixed_attn, the inner checkpoint of the joint attention in layers that are not checkpointed as a whole. - OPENWAM_TEXT_ENCODER_OFFLOAD=1 (with OPENWAM_TEXT_CACHE=1): keep the frozen umT5 weights (10.6 GiB in bf16) in one pinned host buffer per dtype and copy them to the GPU only for an iteration with a text-cache miss. Same weights and kernels, so the rows are bitwise the same.
With the model work reduced, host-side stalls show up in the step time. New switches, all off by default; none of them changes what is computed: - OPENWAM_DS_TIMER_NOSYNC=1: DeepSpeed's throughput timer without its synchronized mode. It calls torch.cuda.synchronize() at every engine forward and step only to time itself, which makes the host wait for the VAE encode before it can queue the first MoT layer. - OPENWAM_DS_OVERLAP_COMM=1: ZeRO-2 overlap_comm, so the gradient bucket all-reduces run on DeepSpeed's reduction stream. - OPENWAM_GC_FREEZE=1: gc.collect() + gc.freeze() after steps 1, 2 and 4 on every rank. The long-lived objects created in the first steps otherwise trigger full collections of about 0.5 s on one rank at a time, which stall all ranks at the next collective. - OPENWAM_NUMA_BIND=1: bind each rank and its DataLoader workers to the CPU cores of its GPU's NUMA node. Both DeepSpeed settings are added to the config built in _build_accelerator; keys set explicitly are kept.
Add two preset files that turn on the speedups for 32 samples per GPU on 8 GPUs (global batch 256, the recipe of the released OpenWAM-Alpha-Sim-LIBERO checkpoint) without gradient accumulation, and let scripts/train.sh load one through OPENWAM_PRESET (unset: no change): - presets/fast.env: all speedups, no activation recompute, for GPUs with more than 80 GB: projected peak 120.4 GiB per GPU on 8 GPUs. - presets/fast_80gb.env: fast.env + umT5 weight offload + recompute of the first 16 MoT layers: projected peak 69.1 GiB per GPU on 8 GPUs, so it fits 80 GB GPUs. On OpenWAM-α LIBERO fine-tuning with 4 GPUs at 32 per GPU, against the default code: 3.13x and 2.41x. presets/README.md lists every switch and whether it keeps training bitwise identical; the README gets a News entry and a short section. A test checks that every switch in a preset is read by the code.
Keep VACE text context intact, require PyTorch 2.7.1+, and fall back to native VAE encoding for single-frame inputs. Preserve wrapped/frozen module forward semantics and text encoder offload lifecycle, and handle fully masked cuDNN attention rows while preserving backward layout. Add targeted regressions. Validation: 181 CPU and 16 CUDA tests passed; three real Inductor block comparisons passed. Four-H200 real-data SFT probes covered seven benchmarks at full and 80 GiB capacity, with additional LIBERO baseline and allocator diagnostics (19 valid runs). These probes are not full convergence evaluations. Documentation changes are intentionally excluded pending user review.
Default training.speed_preset to presets/fast.env, with legacy OPENWAM_PRESET fallback and explicit Hydra path/null overrides. Apply literal preset settings in each worker before training imports and remove early shell sourcing. Validate complete files before changing environment; resolve relative paths from the repository root. Add Hydra, startup-order and shell-launch regressions. Validation: 30 configuration and preset tests pass in local and remote training environments; Ruff and shell syntax checks pass.
Move the first-N layer limit from OPENWAM_CKPT_LAYERS into training.gradient_checkpointing_layers. Default null selects all layers, non-negative integers select the first N, and the existing gradient-checkpointing switch remains required. Preserve the existing dual-system joint-loop scope. Validate layer counts before model loading and in direct driver calls, thread the value through training runtime, and remove both presets layer-count overrides and the obsolete environment parser. Add configuration, runtime, backward-recomputation and gradient-parity coverage. Relevant remote tests: 208 passed, 2 skipped.
wayrise
previously approved these changes
Nov 4, 2026
wayrise
left a comment
wayrise
left a comment
Contributor
There was a problem hiding this comment.
Reviewed commit 4437d3f against main 3950b6d. No blocking regression found in the supported training, resume, and fresh checkpoint-load inference paths.
Rechecked the six correctness fixes, YAML preset/checkpointing/offload defaults, architecture capability guards, gradient accumulation, and README measurement conditions.
Validation on the exact head:
- Repository Ruff, Python compilation, shell syntax, diff whitespace, and train/eval import-boundary checks passed.
- CPU suite: 2,137 passed, 18 skipped, 42 GPU tests deselected. The sole failure, test_cosmos_predict25_build_with_bad_path_raises_clear_error, reproduces with the identical stack on the unchanged base: installed Cosmos/Transformer Engine attempts CUDA initialization while CUDA is hidden. The initial stdin-launch DataLoader failures disappeared with standard python -m pytest.
- With fast.env loaded before imports: 246 architecture/configuration regressions passed, including single/dual/tri and non-Wan paths.
- H200 / PyTorch 2.7.1+cu128: 20 targeted CUDA regressions passed (attention, fused kernels, freezing, VAE, VACE, and text-offload lifecycle).
A non-blocking P3 cache-lifecycle follow-up is recorded inline. The standard entry points configure dtype/device before populating this cache.
The opening PR description still references the removed fast_80gb preset and old opt-in defaults; the current README contains the supported launch/configuration guidance and the qualified four-H200 speedup figures.
GitHub CI and Docker workflows are still action_required with no jobs executed; this review does not treat them as passing.
| ) | ||
| if all(r is not None for r in rows): | ||
| cache.hits += 1 | ||
| return torch.stack(rows, dim=0), seq_lens |
Contributor
There was a problem hiding this comment.
Non-blocking P3: invalidate training text embeddings when an already-used backbone changes dtype/device through set_dtype_device(). This cache is keyed only by token IDs/masks, so a warm-cache hit returns tensors from the previous dtype/device. Reproduced on CPU with the public API: populate a float64 cache, call set_dtype_device(float32, cpu), then encode the same prompt; the cached float64 context fails the float32 text projection, whereas a new prompt succeeds. Clear and recompute the cache on migration (and handle any active offload manager), rather than just casting old embeddings. Standard train/resume/fresh-checkpoint deployment configure dtype/device before cache population, so this does not block those paths.
wayrise
approved these changes
Nov 5, 2026
wayrise
left a comment
wayrise
left a comment
Contributor
There was a problem hiding this comment.
Re-approved the latest head, 66c8d89. The only change since the previous review adds presets/** to the Docker build-context allowlist, fixing the missing default training preset in the image. The existing container regression covers this failure.
Both workflows on this head have now passed:
- CI: https://github.com/OpenWAM-Official/OpenWAM/actions/runs/41958508017
- Docker build and tests: https://github.com/OpenWAM-Official/OpenWAM/actions/runs/41958508172
No new blocking issue found. The previously recorded non-blocking cache-migration follow-up remains applicable.
xueminchi
approved these changes
Nov 5, 2026
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
We use ResearchBoost-Agent to accelerate the OpenWAM-α LIBERO fine-tuning: the agent proposed each change, and a fixed benchmark accepted it only if training stayed bitwise identical to
mainor within its seed-to-seed variation.OPENWAM_*) for the input pipeline, attention and MoTlayer kernels, the Wan2.2 VAE encode, memory, and DeepSpeed/runtime scheduling.
scripts/train.shthroughOPENWAM_PRESET:presets/fast.envpresets/fast_80gb.envCommits
fix: average gradients over accumulation steps under DeepSpeed 0.18.3+openwam/train/utils/grad_accumulation.py; one call in the training loopperf: add opt-in umT5 output cache and faster input pipelineOPENWAM_TEXT_CACHE,OPENWAM_TEXT_MISS_FAST,OPENWAM_CONTEXT_TRUNC,OPENWAM_FAST_INPUT,OPENWAM_VAE_FIRSTperf: add opt-in attention and MoT layer kernelsOPENWAM_SDPA_CUDNN,OPENWAM_FUSED_ROPE,OPENWAM_FAST_CAT,OPENWAM_COMPILE_BLOCKS(openwam/model/kernels/)perf: add opt-in whole-clip Wan2.2 VAE encodeOPENWAM_VAE_FULLCLIP,OPENWAM_COMPILE_VAE,OPENWAM_VAE_CLperf: add opt-in memory switches for the faster settingsOPENWAM_CKPT_LAYERS,OPENWAM_MIXED_ATTN_CKPT,OPENWAM_TEXT_ENCODER_OFFLOADperf: add opt-in scheduling-only runtime settingsOPENWAM_DS_TIMER_NOSYNC,OPENWAM_DS_OVERLAP_COMM,OPENWAM_GC_FREEZE,OPENWAM_NUMA_BINDfeat: add training speed presets and document thempresets/,OPENWAM_PRESETinscripts/train.sh, README sectionpresets/README.mdlists every switch, what it does, and whether it keeps training bitwise identical.Correctness
fast.envandfast_80gb.envswitches reached 99.05% and 99.40% LIBERO success respectively, while the released OpenWAM-Alpha-Sim-LIBERO checkpoint scores 99.2% under the same evaluation, validating the correctness of our acceleration implementation.