From 5f3bb5fd2ad760c01932bd6471afd9fe0aa46917 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Fri, 6 Mar 2026 22:23:52 +0100 Subject: [PATCH 1/7] utility function to flatten tensor --- src/utilities/tensor.cpp | 16 ++++++++++++++++ src/utilities/tensor.h | 2 ++ src/utilities/tensor_container.h | 3 +++ 3 files changed, 21 insertions(+) diff --git a/src/utilities/tensor.cpp b/src/utilities/tensor.cpp index f434b68..347c2e7 100644 --- a/src/utilities/tensor.cpp +++ b/src/utilities/tensor.cpp @@ -119,6 +119,14 @@ TensorShard shard_view(const Tensor& src, int idx, int num) { return TensorShard{shard, idx, num, src.Sizes}; } +Tensor flat_view(const Tensor& src) { + Tensor dst{src}; + dst.Sizes.fill(0); + dst.Sizes[0] = src.nelem(); + dst.Rank = 1; + return dst; +} + void visit(const std::function& func, SimpleTensorContainer& container) { auto cs = container.num_tensors(); for(std::size_t i = 0; i < cs; ++i) { @@ -168,6 +176,14 @@ GenericTensorContainer shard_empty_container(GenericTensorContainer&& c, int wor return std::move(c); } +GenericTensorContainer flattened_view(const GenericTensorContainer& c) { + std::vector flats(c.num_tensors()); + for (std::size_t i = 0; i < c.num_tensors(); ++i) { + flats.at(i) = flat_view(c.get_tensor(i)); + } + return GenericTensorContainer{flats}; +} + GenericTensorContainer shard_view(const GenericTensorContainer& c, int rank, int world) { std::vector shards(c.num_tensors()); for (std::size_t i = 0; i < c.num_tensors(); ++i) { diff --git a/src/utilities/tensor.h b/src/utilities/tensor.h index 7fd2798..200bad7 100644 --- a/src/utilities/tensor.h +++ b/src/utilities/tensor.h @@ -160,4 +160,6 @@ class TensorShard : public Tensor { }; TensorShard shard_view(const Tensor& src, int idx, int num); +Tensor flat_view(const Tensor& src); + #endif //LLMQ_SRC_UTILS_TENSOR_H diff --git a/src/utilities/tensor_container.h b/src/utilities/tensor_container.h index fa89b88..5653a4c 100644 --- a/src/utilities/tensor_container.h +++ b/src/utilities/tensor_container.h @@ -61,6 +61,9 @@ class GenericTensorContainer final : public SimpleTensorContainer { //! are `nullptr`, but sizes have been set up. GenericTensorContainer shard_empty_container(GenericTensorContainer&& c, int world); +//! Flattens all tensors is the container. +GenericTensorContainer flattened_view(const GenericTensorContainer& c); + //! Shards a non-empty tensor container. The returned container's tensors are _views_ into //! the original container's tensors. GenericTensorContainer shard_view(const GenericTensorContainer& c, int rank, int world); From ef34fad650f08e58a3b66f511e5a3c0b696f2079 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Fri, 6 Mar 2026 22:24:37 +0100 Subject: [PATCH 2/7] make FP8-M work also if first dim is not divisible by 512 --- src/training/adamw_optimizer.cpp | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/src/training/adamw_optimizer.cpp b/src/training/adamw_optimizer.cpp index 99889e7..1b672e1 100644 --- a/src/training/adamw_optimizer.cpp +++ b/src/training/adamw_optimizer.cpp @@ -177,17 +177,27 @@ void AdamWStateManager::allocate_state(IModel& model, cudaStream_t stream, EAllo } mBlocksMScales.resize(mConfig.NumLayers); + if(mMType == ETensorDType::FP8_E4M3) { + auto prepare_shape_for_scales = [&](auto&& c) { + // creates shards same as main weight + auto sharded = shard_empty_container(flattened_view(c), mWorld); + // flatten the local shard + auto flattened = flattened_view(sharded); + // and group into scaling groups + auto grouped = shard_empty_container(std::move(flattened), 128); + return grouped; + }; // we "shard" for 128 as many GPUs, so that we get 1 scale per 128 weights. for (int i = 0; i < mConfig.NumLayers; ++i) { - mBlocksMScales[i] = shard_empty_container(model.create_block_container(mConfig, ETensorDType::FP32, ETensorDType::FP32), 128 * mWorld); + mBlocksMScales[i] = prepare_shape_for_scales(model.create_block_container(mConfig, ETensorDType::FP32, ETensorDType::FP32)); alloc_lazy.allocate(mBlocksMScales[i]); alloc_lazy.commit(alloc, EAllocationType::ON_DEVICE, "m_block_scales"); visit([stream](Tensor& t){ fill_constant(t, 1.f, t.nelem(), stream); }, mBlocksMScales[i]); } - mNonBlockMScales = shard_empty_container(model.create_non_block_container(mConfig, ETensorDType::FP32, ETensorDType::FP32), 128 * mWorld); + mNonBlockMScales = prepare_shape_for_scales(model.create_non_block_container(mConfig, ETensorDType::FP32, ETensorDType::FP32)); alloc_lazy.allocate(mNonBlockMScales); alloc_lazy.commit(alloc, EAllocationType::ON_DEVICE, "m_nonblock_scales"); visit([stream](Tensor& t){ From 62512831c658c2ab864bbc097ea412c8e57871e2 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Fri, 6 Mar 2026 22:24:50 +0100 Subject: [PATCH 3/7] prevent crash when trying to fill an empty tensor --- src/kernels/fill.cu | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/kernels/fill.cu b/src/kernels/fill.cu index 8b63e92..a6ac6aa 100644 --- a/src/kernels/fill.cu +++ b/src/kernels/fill.cu @@ -17,6 +17,8 @@ __global__ void fill_kernel(floatX* dst, floatX value, std::size_t count) { template void fill_imp(floatX* dst, floatX value, std::size_t count, cudaStream_t stream) { + if (count == 0) return; + if (dst == nullptr) throw std::invalid_argument("dst is nullptr"); fill_kernel<<(256)), 256, 0, stream>>> (dst, value, count); CUDA_CHECK(cudaGetLastError()); } From 0631202e3da63076753126f33505fa9efa8cd9b1 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Wed, 3 Jun 2026 16:34:47 +0200 Subject: [PATCH 4/7] fixups --- src/kernels/random.cu | 1 + src/training/adamw_optimizer.cpp | 8 ++++---- src/utilities/tensor.cpp | 4 ++-- src/utilities/tensor_container.h | 2 +- 4 files changed, 8 insertions(+), 7 deletions(-) diff --git a/src/kernels/random.cu b/src/kernels/random.cu index ff29b3d..5e6ff3d 100644 --- a/src/kernels/random.cu +++ b/src/kernels/random.cu @@ -28,6 +28,7 @@ __global__ void rng_normal_kernel(floatX* dst, std::size_t count, float mean, fl template void rng_normal_imp(floatX* dst, std::size_t count, float mean, float std, unsigned long long seed, unsigned long long subsequence, cudaStream_t stream) { assert(count % 4 == 0); + if (count == 0) return; rng_normal_kernel<<(4*256)), 256, 0, stream>>> (dst, count, mean, std, seed, subsequence); CUDA_CHECK(cudaGetLastError()); } diff --git a/src/training/adamw_optimizer.cpp b/src/training/adamw_optimizer.cpp index 1b672e1..de124c0 100644 --- a/src/training/adamw_optimizer.cpp +++ b/src/training/adamw_optimizer.cpp @@ -182,13 +182,13 @@ void AdamWStateManager::allocate_state(IModel& model, cudaStream_t stream, EAllo auto prepare_shape_for_scales = [&](auto&& c) { // creates shards same as main weight auto sharded = shard_empty_container(flattened_view(c), mWorld); - // flatten the local shard - auto flattened = flattened_view(sharded); // and group into scaling groups - auto grouped = shard_empty_container(std::move(flattened), 128); + auto grouped = shard_empty_container(std::move(sharded), 128); return grouped; }; - // we "shard" for 128 as many GPUs, so that we get 1 scale per 128 weights. + // we first shard by mWorld (matching main weights), then shard the local + // flattened view by 128 to get 1 scale per 128 weights. + for (int i = 0; i < mConfig.NumLayers; ++i) { mBlocksMScales[i] = prepare_shape_for_scales(model.create_block_container(mConfig, ETensorDType::FP32, ETensorDType::FP32)); alloc_lazy.allocate(mBlocksMScales[i]); diff --git a/src/utilities/tensor.cpp b/src/utilities/tensor.cpp index 347c2e7..dc47180 100644 --- a/src/utilities/tensor.cpp +++ b/src/utilities/tensor.cpp @@ -121,7 +121,7 @@ TensorShard shard_view(const Tensor& src, int idx, int num) { Tensor flat_view(const Tensor& src) { Tensor dst{src}; - dst.Sizes.fill(0); + dst.Sizes.fill(1); dst.Sizes[0] = src.nelem(); dst.Rank = 1; return dst; @@ -181,7 +181,7 @@ GenericTensorContainer flattened_view(const GenericTensorContainer& c) { for (std::size_t i = 0; i < c.num_tensors(); ++i) { flats.at(i) = flat_view(c.get_tensor(i)); } - return GenericTensorContainer{flats}; + return GenericTensorContainer{std::move(flats)}; } GenericTensorContainer shard_view(const GenericTensorContainer& c, int rank, int world) { diff --git a/src/utilities/tensor_container.h b/src/utilities/tensor_container.h index 5653a4c..1e5f1a4 100644 --- a/src/utilities/tensor_container.h +++ b/src/utilities/tensor_container.h @@ -61,7 +61,7 @@ class GenericTensorContainer final : public SimpleTensorContainer { //! are `nullptr`, but sizes have been set up. GenericTensorContainer shard_empty_container(GenericTensorContainer&& c, int world); -//! Flattens all tensors is the container. +//! Flattens all tensors in the container. GenericTensorContainer flattened_view(const GenericTensorContainer& c); //! Shards a non-empty tensor container. The returned container's tensors are _views_ into From 0241e4ea69ec1367518dcb414fcb017e955a9806 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sat, 18 Jul 2026 14:20:39 +0200 Subject: [PATCH 5/7] Fix disabled-tensor handling: null sentinels, zero-nelem rank-0 tensors, 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 --- src/models/llama_model.cpp | 2 +- src/training/gradients.cpp | 6 ++++-- src/utilities/lazy_allocator.cpp | 5 +++++ src/utilities/tensor.cpp | 3 +++ src/utilities/tensor.h | 5 +++++ 5 files changed, 18 insertions(+), 3 deletions(-) diff --git a/src/models/llama_model.cpp b/src/models/llama_model.cpp index c4e704d..80b366c 100644 --- a/src/models/llama_model.cpp +++ b/src/models/llama_model.cpp @@ -776,7 +776,7 @@ void LLamaModel::_calculate_gradient_norm(NCCLCommunicator& comm, float grad_cli auto& rs = RunState; fill_zero(rs->NormBuffer, stream); - auto norm_squared = [&](const TensorShard& grad){ + auto norm_squared = [&](const Tensor& grad){ global_norm_squared(rs->NormBuffer, grad, grad.nelem(), rs->DeviceProp, stream); }; diff --git a/src/training/gradients.cpp b/src/training/gradients.cpp index 55f5744..253782d 100644 --- a/src/training/gradients.cpp +++ b/src/training/gradients.cpp @@ -167,8 +167,10 @@ SimpleTensorContainer& ShardedBlocksGradientManager::get_block_full(int layer_id void ShardedBlocksGradientManager::notify_non_block(std::size_t index, cudaStream_t stream, NCCLCommunicator& comm) { if(!is_last_micro_step()) return; - NvtxRange r{"notify"}; - comm.reduce_scatter(mFullNonBlock.get_tensor(index), stream, mNonBlockEvent); + if (comm.world_size() != 1) { + NvtxRange r{"notify"}; + comm.reduce_scatter(mFullNonBlock.get_tensor(index), stream, mNonBlockEvent); + } } void ShardedBlocksGradientManager::notify_block(int layer_idx, cudaStream_t stream, NCCLCommunicator& comm) { diff --git a/src/utilities/lazy_allocator.cpp b/src/utilities/lazy_allocator.cpp index 3241eee..24230ee 100644 --- a/src/utilities/lazy_allocator.cpp +++ b/src/utilities/lazy_allocator.cpp @@ -39,6 +39,9 @@ Tensor LazyAllocator::commit(TensorAllocator& storage, EAllocationType type, con Tensor backing = storage.allocate(ETensorDType::BYTE, name, type, {(long)total_size}); auto* ptr = backing.get(); for(auto& target: mTargets) { + // zero-size tensors keep Data == nullptr; a null Data pointer is the sentinel for + // "disabled" tensors (e.g., QKV bias, tied LM head) that `visit` and other checks rely on + if(target->bytes() == 0) continue; target->Data = ptr; target->Device = backing.Device; ptr += div_ceil(target->bytes(), page_size) * page_size; @@ -63,6 +66,8 @@ Tensor LazyAllocator::commit(DeviceMemoryStack& storage, const char* name) { if (backing) { auto* ptr = backing.get(); for(auto& target: mTargets) { + // zero-size tensors keep Data == nullptr, see above + if(target->bytes() == 0) continue; target->Data = ptr; target->Device = backing.Device; ptr += div_ceil(target->bytes(), page_size) * page_size; diff --git a/src/utilities/tensor.cpp b/src/utilities/tensor.cpp index dc47180..700c7c5 100644 --- a/src/utilities/tensor.cpp +++ b/src/utilities/tensor.cpp @@ -100,6 +100,9 @@ TensorShard::TensorShard(const Tensor& src) : Tensor(src), GlobalShape(src.Sizes } std::size_t TensorShard::global_nelem() const { + // rank 0 denotes a disabled tensor holding no elements, same as Tensor::nelem() + if (Rank == 0) + return 0; std::size_t sz = 1; for (int i = 0; i < Rank; ++i) sz *= GlobalShape[i]; diff --git a/src/utilities/tensor.h b/src/utilities/tensor.h index 200bad7..803bb4f 100644 --- a/src/utilities/tensor.h +++ b/src/utilities/tensor.h @@ -37,6 +37,11 @@ struct Tensor { } [[nodiscard]] constexpr std::size_t nelem() const { + // rank 0 denotes a default-constructed/disabled tensor, not a scalar, + // so it holds no elements (rather than the empty product 1) + if(Rank == 0) { + return 0; + } std::size_t sz = 1; for(int i = 0; i < Rank; ++i) { sz *= Sizes[i]; From 810d0e4a1065149a24f922345fac269ef94d411d Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sat, 18 Jul 2026 18:59:12 +0200 Subject: [PATCH 6/7] keep momentum of 1D tensors unquantized under FP8-M 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 --- src/models/llama_optimizer.cpp | 33 ++++++++++++++++---------------- src/training/adamw_optimizer.cpp | 31 +++++++++++++++++++----------- src/training/adamw_optimizer.h | 1 + 3 files changed, 37 insertions(+), 28 deletions(-) diff --git a/src/models/llama_optimizer.cpp b/src/models/llama_optimizer.cpp index bcad5e8..52a1f8e 100644 --- a/src/models/llama_optimizer.cpp +++ b/src/models/llama_optimizer.cpp @@ -23,27 +23,26 @@ struct OptStateWrapper : ITensorContainer { }; void OptStateWrapper::iterate_tensors(const std::function& callback) { - callback("model.embed_tokens.weight", NonBlock->get_tensor(LLamaWeightID::EMBEDDING)); - if(NonBlock->get_tensor(LLamaWeightID::LM_HEAD)) { - callback("lm_head.weight", NonBlock->get_tensor(LLamaWeightID::LM_HEAD)); - } - callback("model.norm.weight", NonBlock->get_tensor(LLamaWeightID::LNF_W)); + auto cb = [&callback](std::string name, const Tensor& t) { + if (t) { + callback(std::move(name), t); + } + }; + + cb("model.embed_tokens.weight", NonBlock->get_tensor(LLamaWeightID::EMBEDDING)); + cb("lm_head.weight", NonBlock->get_tensor(LLamaWeightID::LM_HEAD)); + cb("model.norm.weight", NonBlock->get_tensor(LLamaWeightID::LNF_W)); for(int i = 0; i < Blocks->size(); i++) { auto& layer = Blocks->at(i); - const Tensor& qkv_w = layer.get_tensor(LLamaWeightID::QKV_W); - const Tensor& up_proj = layer.get_tensor(LLamaWeightID::UP_W); std::string prefix = "model.layers." + std::to_string(i); - callback(prefix + ".self_attn.qkv.weight", qkv_w); - if (layer.get_tensor(LLamaWeightID::QKV_B)) { - callback(prefix + ".self_attn.qkv.bias", layer.get_tensor(LLamaWeightID::QKV_B)); - } - - callback(prefix + ".self_attn.o_proj.weight", layer.get_tensor(LLamaWeightID::ATTO_W)); - callback(prefix + ".mlp.up.weight", up_proj); - callback(prefix + ".mlp.down_proj.weight", layer.get_tensor(LLamaWeightID::DOWN_W)); - callback(prefix + ".input_layernorm.weight", layer.get_tensor(LLamaWeightID::LN1_W)); - callback(prefix + ".post_attention_layernorm.weight", layer.get_tensor(LLamaWeightID::LN2_W)); + cb(prefix + ".self_attn.qkv.weight", layer.get_tensor(LLamaWeightID::QKV_W)); + cb(prefix + ".self_attn.qkv.bias", layer.get_tensor(LLamaWeightID::QKV_B)); + cb(prefix + ".self_attn.o_proj.weight", layer.get_tensor(LLamaWeightID::ATTO_W)); + cb(prefix + ".mlp.up.weight", layer.get_tensor(LLamaWeightID::UP_W)); + cb(prefix + ".mlp.down_proj.weight", layer.get_tensor(LLamaWeightID::DOWN_W)); + cb(prefix + ".input_layernorm.weight", layer.get_tensor(LLamaWeightID::LN1_W)); + cb(prefix + ".post_attention_layernorm.weight", layer.get_tensor(LLamaWeightID::LN2_W)); } } diff --git a/src/training/adamw_optimizer.cpp b/src/training/adamw_optimizer.cpp index de124c0..75f9083 100644 --- a/src/training/adamw_optimizer.cpp +++ b/src/training/adamw_optimizer.cpp @@ -18,8 +18,8 @@ AdamWStateManager::AdamWStateManager(TransformerConfig cfg, IModel& model, bool mConfig(cfg), mOffloadM(offload_m), mOffloadV(offload_v), mUseZeroCopy(zero_copy), mRank(rank), mWorld(world), mMType(type_m), mVType(type_v) { if(mOffloadM && !mUseZeroCopy) { - mMDeviceBuffer[0] = shard_empty_container(model.create_block_container(mConfig, mMType, mMType), mWorld); - mMDeviceBuffer[1] = shard_empty_container(model.create_block_container(mConfig, mMType, mMType), mWorld); + mMDeviceBuffer[0] = shard_empty_container(model.create_block_container(mConfig, mMType, non_matrix_m_type()), mWorld); + mMDeviceBuffer[1] = shard_empty_container(model.create_block_container(mConfig, mMType, non_matrix_m_type()), mWorld); } if(mOffloadV && !mUseZeroCopy) { @@ -158,17 +158,21 @@ void AdamWStateManager::store_block(int layer_idx, cudaStream_t stream, cudaStre } } +ETensorDType AdamWStateManager::non_matrix_m_type() const { + return mMType == ETensorDType::FP8_E4M3 ? ETensorDType::BF16 : mMType; +} + void AdamWStateManager::allocate_state(IModel& model, cudaStream_t stream, EAllocationType kind, TensorAllocator& alloc) { { auto ctx = alloc.with_context("Adam M"); LazyAllocator alloc_lazy; mBlocksM.resize(mConfig.NumLayers); for (int i = 0; i < mConfig.NumLayers; ++i) { - mBlocksM[i] = shard_empty_container(model.create_block_container(mConfig, mMType, mMType), mWorld); + mBlocksM[i] = shard_empty_container(model.create_block_container(mConfig, mMType, non_matrix_m_type()), mWorld); alloc_lazy.allocate(mBlocksM[i]); mStorageM.push_back(alloc_lazy.commit(alloc, mOffloadM ? kind : EAllocationType::ON_DEVICE, "m_block_shard")); } - mNonBlockM = shard_empty_container(model.create_non_block_container(mConfig, mMType, mMType), mWorld); + mNonBlockM = shard_empty_container(model.create_non_block_container(mConfig, mMType, non_matrix_m_type()), mWorld); alloc_lazy.allocate(mNonBlockM); mStorageM.push_back(alloc_lazy.commit(alloc, mOffloadM ? kind : EAllocationType::ON_DEVICE, "m_nonblock_shard")); @@ -179,16 +183,21 @@ void AdamWStateManager::allocate_state(IModel& model, cudaStream_t stream, EAllo mBlocksMScales.resize(mConfig.NumLayers); if(mMType == ETensorDType::FP8_E4M3) { - auto prepare_shape_for_scales = [&](auto&& c) { - // creates shards same as main weight - auto sharded = shard_empty_container(flattened_view(c), mWorld); - // and group into scaling groups - auto grouped = shard_empty_container(std::move(sharded), 128); - return grouped; - }; // we first shard by mWorld (matching main weights), then shard the local // flattened view by 128 to get 1 scale per 128 weights. + auto prepare_shape_for_scales = [&](GenericTensorContainer&& c) { + for (std::size_t i = 0; i < c.num_tensors(); ++i) { + auto& t = c.get_tensor(i); + // only apply to 2D weights + if (t.Rank != 2) { + t.Rank = 1; + t.Sizes[0] = 0; + } + } + return shard_empty_container(shard_empty_container(flattened_view(c), mWorld), 128); + }; + for (int i = 0; i < mConfig.NumLayers; ++i) { mBlocksMScales[i] = prepare_shape_for_scales(model.create_block_container(mConfig, ETensorDType::FP32, ETensorDType::FP32)); alloc_lazy.allocate(mBlocksMScales[i]); diff --git a/src/training/adamw_optimizer.h b/src/training/adamw_optimizer.h index 9535c87..95a44f0 100644 --- a/src/training/adamw_optimizer.h +++ b/src/training/adamw_optimizer.h @@ -41,6 +41,7 @@ class AdamWStateManager { protected: SimpleTensorContainer& get_block_from(int layer_idx, cudaStream_t stream, SimpleTensorContainer& buf); + [[nodiscard]] ETensorDType non_matrix_m_type() const; TransformerConfig mConfig; bool mOffloadM; From 3edf49f635cf0ae7ff995f454e986318fb5f6706 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sat, 18 Jul 2026 23:11:31 +0200 Subject: [PATCH 7/7] don't crash when the GPU monitoring thread cannot set CPU affinity MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- src/utilities/gpu_info_nvml.cpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/utilities/gpu_info_nvml.cpp b/src/utilities/gpu_info_nvml.cpp index b675772..e324bc4 100644 --- a/src/utilities/gpu_info_nvml.cpp +++ b/src/utilities/gpu_info_nvml.cpp @@ -128,7 +128,11 @@ void GPUUtilTrackerNVML::setup_tracking_thread() { // TODO should this be one thread for all devices? mThread = std::jthread([this](std::stop_token stop_token) { - NVML_CHECK(nvmlDeviceSetCpuAffinity(mDevice)); + // best-effort: sandboxed environments (e.g. gVisor on Modal) reject + // affinity operations, and a monitoring thread must not kill training + if (nvmlDeviceSetCpuAffinity(mDevice) != NVML_SUCCESS) { + fprintf(stderr, "[NVML WARNING] could not set CPU affinity for GPU monitoring thread\n"); + } nvmlFieldValue_t fields[] = {{NVML_FI_DEV_PCIE_COUNT_RX_BYTES}, {NVML_FI_DEV_PCIE_COUNT_TX_BYTES}, {NVML_FI_DEV_TOTAL_ENERGY_CONSUMPTION, 0}}; while (true) {