Repository navigation
Speed up OpenWAM-α fine-tuning 2.4-3.1x with ResearchBoost-Agent - #58
Open
GONGSHUKAI wants to merge 7 commits into
Open
GONGSHUKAI wants to merge 7 commits into
GONGSHUKAI wants to merge 7 commits into
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.
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.