Fp8 opt bugfix#68
Merged
Merged
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR fixes a bug in FP8 optimizer state allocation that caused failures when model weight tensors had a first dimension not divisible by 128 × world_size. The fix introduces a flat_view / flattened_view utility to correctly shape FP8 momentum scaling tensors by flattening each weight tensor before sharding, rather than directly dividing the first dimension of multi-dimensional tensors.
Changes:
- New
flat_view(Tensor)andflattened_view(GenericTensorContainer)utility functions to produce 1D views of tensors and containers - Rewrote FP8 scale tensor shaping in
AdamWStateManager::allocate_stateto flatten first, shard bymWorld, then shard by 128 - Added early-exit and null-pointer guards to
fill_impto safely handle zero-element tensors
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
src/utilities/tensor.h |
Declares new flat_view(const Tensor&) function |
src/utilities/tensor.cpp |
Implements flat_view and flattened_view(GenericTensorContainer) |
src/utilities/tensor_container.h |
Declares flattened_view(GenericTensorContainer) |
src/training/adamw_optimizer.cpp |
Replaces direct shard_empty_container(_, 128*mWorld) with a correct flatten-then-shard approach for FP8 scale tensors |
src/kernels/fill.cu |
Adds count == 0 early return and null pointer guard to fill_imp |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
…rs, grad-norm shard type LazyAllocator::commit assigned every registered tensor a non-null Data pointer, including zero-size ones (disabled QKV bias, tied LM head). All 'disabled tensor' checks key off Data == nullptr, so those tensors were suddenly considered live: backward_bias corrupted the attention-out gradient on no-bias models, zero-element kernel launches and NCCL calls crashed sharded/FP8-momentum configs, tied models lost the skip-zeroing fast path, and optimizer-state checkpoints gained a bogus lm_head entry. Commit now leaves zero-size tensors with a null Data pointer. Also: - rank-0 tensors now report nelem() == 0 instead of the empty product 1: the codebase uses default-constructed rank-0 tensors as 'no tensor' (never as scalars), and the phantom element made shard_view's consistency check throw on the disabled QKV bias in shard_block, aborting any multi-GPU run of a no-bias model during weight allocation - ShardedBlocksGradientManager::notify_non_block skips the redundant reduce_scatter on single-GPU runs, matching the unsharded manager - _calculate_gradient_norm takes gradient shards as const Tensor& to avoid synthesizing TensorShard temporaries with wrong metadata Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
With --opt-m-dtype=e4m3, the momentum of every tensor was stored in FP8,
including tiny 1D tensors (norm weights, qkv bias). Their scale tensors
required the per-rank shard size to be divisible by 128, which fails for
e.g. hidden size 896 on 2 or 4 GPUs ("Cannot divide 448 by 128").
Quantizing these tensors saves almost no memory, so store their momentum
in BF16 instead and give them empty scale entries. Empty tensors are now
skipped uniformly in OptStateWrapper::iterate_tensors, which also covers
the tied lm_head and disabled qkv bias cases.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
nvmlDeviceSetCpuAffinity fails with NVML_ERROR_UNKNOWN in sandboxed environments (e.g. gVisor on Modal). The NVML_CHECK threw inside the monitoring jthread, where the uncaught exception called std::terminate and took down the whole training process — on Modal CI this killed the worker mid-input, surfacing as "Server has lost track of input". Affinity for the monitoring thread is an optimization; warn and continue, matching how set_cpu_affinity failures are handled in comm.cpp. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 12 out of 12 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (1)
src/training/gradients.cpp:173
- When
comm.world_size() == 1, this function skipsreduce_scatterand never recordsmNonBlockEvent. However,end_micro_step()unconditionally callscudaStreamWaitEvent(stream, mNonBlockEvent, 0)on the last micro-step, which can block indefinitely on an event that was never recorded.
void ShardedBlocksGradientManager::notify_non_block(std::size_t index, cudaStream_t stream, NCCLCommunicator& comm) {
if(!is_last_micro_step()) return;
if (comm.world_size() != 1) {
NvtxRange r{"notify"};
comm.reduce_scatter(mFullNonBlock.get_tensor(index), stream, mNonBlockEvent);
}
Comment on lines
+64
to
+66
| //! Flattens all tensors in the container. | ||
| GenericTensorContainer flattened_view(const GenericTensorContainer& c); | ||
|
|
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.
seems to work again for 1.5b