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()); } 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/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/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 99889e7..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")); @@ -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){ 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; 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/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) { 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 f434b68..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]; @@ -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& 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 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 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..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]; @@ -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 diff --git a/src/utilities/tensor_container.h b/src/utilities/tensor_container.h index fa89b88..1e5f1a4 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 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);