Add TAO-compatible Cosmos3 fine-tuning support - #164
Draft
ramanathan831 wants to merge 40 commits into
Draft
Conversation
TAO video recipes resolve the dataloader seed via ${oc.env:TAO_DATALOADER_SEED,42},
which OmegaConf returns as a string; stream() then computed seed + epoch and
raised 'can only concatenate str (not int) to str' on the first batch.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
get_rope_index built per-sample position_ids shaped (3, L_i) and passed them straight to pad_sequence, which requires the variable dimension first; any generation batch mixing sequence lengths crashed with a size-mismatch error. Transpose to (L_i, 3) before padding and permute back to (3, batch, L_max). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ity' into dev/ram/cosmos3-tao-reproducibility # Conflicts: # cosmos_framework/data/generator/dataflow/distributors.py # cosmos_framework/data/generator/dataflow/resume_test.py
The Edge reasoner inherited the Qwen3-VL-specific 'cosmos' NATTEN adapter, which rejects Edge's explicit attention mask, so Edge never trained without a manual attn_implementation override — and the profile default flash_attention_2 has no aarch64 wheels. Default the Edge policy to flash_attention_2 when flash_attn is importable and sdpa otherwise, so a fresh build runs on both x86 and ARM without spec-level overrides. Explicit TOML values still win. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Ramanathan Arunachalam <rarunachalam@nvidia.com>
init_flash_attn_meta(deterministic=True) set cudnn.deterministic, use_deterministic_algorithms and CUBLAS_WORKSPACE_CONFIG, and its docstring claimed "HF flash_attention_2 respects torch.backends.cudnn.deterministic and torch.use_deterministic_algorithms()". That is not the case. FlashAttention-2 is a separate CUDA extension whose backward uses atomics, so use_deterministic_algorithms cannot police it -- with warn_only=False it reports nothing at all and the run is still nondeterministic. Transformers selects FA2's deterministic backward from FLASH_ATTENTION_DETERMINISTIC in modeling_flash_attention_utils, which nothing set. Measured on the sibling Cosmos-RL backend, which had the identical gap: two same-seed runs matched for two steps and then diverged, and setting only CUBLAS_WORKSPACE_CONFIG or only warn_only=False did not help. With this variable set, two full 540-step runs were bit-identical in loss and gradient norm at a cost of 1.60s/step against 1.62s. Signed-off-by: Ramanathan Arunachalam <rarunachalam@nvidia.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
precision reaches the network only through MixedPrecisionPolicy, which fully_shard attaches. The dp_shard <= 1 path returns before that, so a single-GPU run silently keeps whatever dtype the checkpoint loaded as. For a Qwen3-VL base that is fp32, and every matmul then dispatches to SIMT fp32 CUTLASS kernels instead of tensor cores. Measured on a GB300 with Cosmos3-Nano LoRA SFT: fp32 aten::mm was 89% of training CUDA time, and the same GEMM takes 3.97ms in fp32 against 0.14ms in bf16. Multi-GPU runs never showed this because dp_shard > 1 attaches the policy, which is why the two backends were at parity on an 8-GPU node and 3.8x apart on one GPU. Cast the frozen parameters to the configured compute dtype on that path. Trainable parameters are left alone so optimizer state keeps the precision it was built with, mirroring what MixedPrecisionPolicy does for master weights. Steady state on one GB300, batch 31: 5.86s/step -> 3.42s/step (1.71x), epoch 17.6min -> 10.3min, GPU memory 213GB -> 186GB, step-1 loss 0.8710 against 0.8759 (bf16 rounding). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Ramanathan Arunachalam <rarunachalam@nvidia.com>
The backend runs through PyTorch's cuDNN SDPA ATen op rather than the standalone cuDNN frontend, and that op is present and correct from 9.15. Requiring 9.20 excluded the path entirely on stacks shipping 9.15.1. That exclusion is not cosmetic on Blackwell. The arch-103 preference order is cudnn, natten, flash2; flash2 needs a flash-attn package the CUDA base image does not carry, so declining cuDNN leaves NATTEN as the only usable backend. Measured on a GB300 with cuDNN 91501. Against the math reference the cuDNN kernel gives max|err| 0.01562, identical to the FLASH and EFFICIENT SDPA backends, and it is the fastest of them (1.43ms vs 2.75ms for flash on a 31x704x32x128 forward+backward). End to end on Cosmos3-Nano LoRA SFT it takes 2.94s/step to 2.62s/step with training loss preserved (0.8664 against 0.8759 for the fp32 reference, within bf16 rounding). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Ramanathan Arunachalam <rarunachalam@nvidia.com>
LoRA SFT at a constant learning rate with no decay is marginally stable. Most steps sit well under the clip threshold, then one unlucky batch produces a gradient norm one or two orders of magnitude larger. Clipping bounds that step's magnitude but not its direction, and AdamW folds the anomalous direction into its moments; because its update is scale-invariant, the poisoned moments keep steering the model long after the batch is gone. Two unguarded runs of the same recipe on a GB300, differing only in attention backend, both diverged -- one at step 188 (recovered, MCQ accuracy 0.8568), one at step 466 (never recovered, final validation loss 0.766, accuracy 0.1905). The step is not reproducible run to run, so this is a stochastic property of the recipe rather than of any one backend. Detection uses the gradient norm because the loss lags: at the spike the norm separates from its running median by 10.4x while the loss still reads 0.2742 and separates by only 2.8x. Recovery rewinds rather than skips. Skipping leaves the poisoned moments in place, and once the loss is elevated every subsequent norm looks like a spike, so a skip-based guard freezes a diverged run instead of rescuing it. The guard keeps a short ring of parameter and optimizer-moment snapshots, restores the oldest, and backs the learning rate off. Two details specific to this codebase. Detection runs in on_after_backward because GradClip rescales every spike down to clip_norm, leaving no signal in a post-clip norm. The backoff rescales scheduler.base_lrs rather than param_group["lr"], because LambdaLR recomputes lr = base_lrs * lambda(step) on every step and would silently discard a direct write. Default off: the snapshot ring scales with the trainable parameter count, roughly 1.2GB for a rank-64 LoRA adapter but hundreds of GB for dense 8B fine-tuning. Measured overhead with the ring active: 2.62 to 2.73 s/step, 4.2%. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Ramanathan Arunachalam <rarunachalam@nvidia.com>
A run escaped the guard despite six rollbacks, ending at validation loss 0.858 and MCQ accuracy 0.1619. Its own log shows why: at iteration 301 the guard reported "gradient norm 311.24 exceeds 10.0x the median of the last 47 steps (9.0293)". A healthy median for this recipe is about 0.3 to 0.8, so the baseline had inflated more than an order of magnitude and firing now required a norm above 90. The guard went quiet exactly when it was needed. Three causes, all fixed here. The baseline was purely relative, so it adapted to sickness as though it were the new normal. It is now anchored: it may not exceed baseline_inflation_cap times the lowest median the run has demonstrated. The gradient-norm window was cleared on every rollback, discarding the only healthy reference available and forcing the baseline to be rebuilt from post-rollback steps that were themselves unhealthy. The clear was never needed, because a spiking value is returned before it is ever appended -- the window only ever held clean samples. Learning-rate recovery climbed back to full. At 1.02 per clean step a rate backed off to 0.305 returns to roughly 0.82 within fifty steps, i.e. back to a rate the run had already proven it could not hold. The ceiling now ratchets down by lr_ceiling_decay on each rollback and recovery targets that ceiling. Evidence: a run under the fixed guard absorbed a spike of essentially the same magnitude as the one that destroyed the run above (6.11 against 6.15), finishing at validation loss 0.221 with the lowest epoch-3 loss mean of any run measured, 0.055. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Ramanathan Arunachalam <rarunachalam@nvidia.com>
The guard rescued a run from divergence and then undertrained it. Four rollbacks inside thirteen steps compounded 0.5**4 straight into the rate floor, and the run spent its remaining 274 of 540 steps at a tenth of the intended rate. Maximum loss stayed healthy at 0.97, but MCQ accuracy came out at 0.7858 against roughly 0.88 for runs that were not suppressed -- divergence traded for undertraining. Two changes. A burst of spikes is now one episode rather than one escalation per rollback: only the first rollback within backoff_cooldown steps moves the rate. Rewinding still happens every time, because rewinding is cheap and safe while cutting the rate repeatedly is not. And the backoff is measured from the ceiling instead of from the current scale, so the depth of a dip no longer depends on how many dips preceded it; the ceiling alone carries the persistent penalty. Replaying that run's fourteen rollback iterations through both policies: the old one reaches the floor at step 266 and stays there for eleven of the fourteen, while the new one never floors, bottoming at 0.205 while still ratcheting 0.4, 0.32, 0.256, 0.205 as trouble recurs. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Ramanathan Arunachalam <rarunachalam@nvidia.com>
Recovery was gated on "this step did not spike", which a degraded run satisfies most of the time. One run collapsed at step 253 and sat at a median loss of 0.69, roughly five times its healthy value, for the remaining 280 steps -- and the rate walked back up to 0.627 of base while it did, because the steps between its spikes looked clean. Recovery now additionally requires the median gradient norm to be within recovery_health_factor of the lowest median the run has demonstrated. The guard already tracked that figure for the anti-inflation cap; this reuses it as evidence of recovery rather than only as a detection floor. A few recovery steps still land before the window accumulates enough elevated samples to register the degradation. That is intended -- the alternative is refusing to recover on transient noise -- and recovery then stops well short of the ceiling an ungated walk would have reached. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Ramanathan Arunachalam <rarunachalam@nvidia.com>
processed_video_cache_size was pinned to 0 whenever shuffle was set, so only validation ever got a cache and training re-decoded all 5,555 videos on every epoch. Decode, not compute, is what a training step here costs. Timing the loop directly on a GB300 splits the steady-state step into 2.00s waiting on the dataloader against 0.62s of compute -- roughly three quarters of the run spent waiting, with the GPU idle for most of it. Enabling the cache takes the wait to 0.02s and the step from 2.62s to 0.69s. Warm epochs drop from 7.86 to 2.57 minutes of training. The gain appears even in the first epoch because the dataset carries several records per video, so repeats hit the cache immediately. Numerically inert, as expected of a decoded-frame cache: same seed, same three validation losses to three decimals, same MCQ accuracy of 0.8929. Defaults to 0. The cache is per dataloader worker and holds decoded frames, so capacity has to be chosen against dataset size, worker count and host memory rather than assumed; 6000 entries for this dataset was comfortable in 736GB. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Ramanathan Arunachalam <rarunachalam@nvidia.com>
Upstream Qwen3VLVisionAttention reserves its packed path for FlashAttention-2, which takes cu_seqlens and issues a single varlen kernel. Every other backend falls through to a loop that splits q/k/v on cu_seqlens and calls attention once per chunk. There is no aarch64 flash-attn wheel for this CUDA base image, so the loop is what runs: 124 chunks x 27 blocks = 3,348 calls per step, each of shape [1, 16, 32, 72], costing 161us of CPU dispatch for 3.6us of GPU work. When every chunk is the same length -- true whenever the batch is uniform in resolution and frame count -- the loop is exactly a batched attention over the chunk axis, so the chunks fold into dim 0 and issue as one call. This is an identity, not an approximation: attention never mixes across chunks in either form and each chunk keeps its own softmax normalisation. Verified at the real shapes, max|difference| 0.000e+00. Ragged batches keep the stock loop. Honest accounting of the payoff: this removes most of the host-side dispatch, taking TorchDynamo cache lookups from 14,004 calls (2.14s) to 621 (128ms), but wall time only improves about 2.6%. Dispatch was not the critical path -- video decode was, and is addressed separately. Kept because it is free and correct, and because it removes the overhead that would otherwise dominate once decode is cached. Behind TAO_FRAMEWORK_BATCH_VISION_ATTENTION, default on. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Ramanathan Arunachalam <rarunachalam@nvidia.com>
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
Add repository-owned Cosmos3-Nano and Edge training compatibility, dataset/checkpoint preparation, PEFT support, status callbacks, validation metrics, and export/evaluation integration.
This is the GitHub-main replacement for work developed before the GitLab repositories were discontinued.
Migration contract
maindev/ram/cosmos3-tao-reproducibilitydev/ram/