Fp8 opt bugfix

src/kernels/fill.cu

@@ -17,6 +17,8 @@ __global__ void fill_kernel(floatX* dst, floatX value, std::size_t count) {
 
 template<typename floatX>
 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<<<div_ceil(count, static_cast<std::size_t>(256)), 256, 0, stream>>> (dst, value, count);
     CUDA_CHECK(cudaGetLastError());
 }

src/kernels/random.cu

@@ -28,6 +28,7 @@ __global__ void rng_normal_kernel(floatX* dst, std::size_t count, float mean, fl
 template<typename floatX>
 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<<<div_ceil(count, static_cast<std::size_t>(4*256)), 256, 0, stream>>> (dst, count, mean, std, seed, subsequence);
     CUDA_CHECK(cudaGetLastError());
 }

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);
     };
 

src/models/llama_optimizer.cpp

@@ -23,27 +23,26 @@ struct OptStateWrapper : ITensorContainer {
 };
 
 void OptStateWrapper::iterate_tensors(const std::function<void(std::string, const TensorShard&)>& 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));
     }
 }
 

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"));
 
@@ -177,17 +181,32 @@ void AdamWStateManager::allocate_state(IModel& model, cudaStream_t stream, EAllo
         }
 
         mBlocksMScales.resize(mConfig.NumLayers);
+
         if(mMType == ETensorDType::FP8_E4M3) {
-            // 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.
+
+            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] = 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){

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;

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) {

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) {

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<std::byte>();
     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<std::byte>();
         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;

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];
@@ -119,6 +122,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(1);
+    dst.Sizes[0] = src.nelem();
+    dst.Rank = 1;
+    return dst;
+}
+
 void visit(const std::function<void(Tensor&)>& func, SimpleTensorContainer& container) {
     auto cs = container.num_tensors();
     for(std::size_t i = 0; i < cs; ++i) {
@@ -168,6 +179,14 @@ GenericTensorContainer shard_empty_container(GenericTensorContainer&& c, int wor
     return std::move(c);
 }
 
+GenericTensorContainer flattened_view(const GenericTensorContainer& c) {
+    std::vector<Tensor> 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{std::move(flats)};
+}
+
 GenericTensorContainer shard_view(const GenericTensorContainer& c, int rank, int world) {
     std::vector<Tensor> shards(c.num_tensors());
     for (std::size_t i = 0; i < c.num_tensors(); ++i) {

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];
@@ -160,4 +165,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

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 in 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);