From e4b3fa099f756f2e5135ef79f21a7f4da88ac2bf Mon Sep 17 00:00:00 2001 From: zhengjiahong Date: Mon, 14 Sep 2026 09:59:53 +0800 Subject: [PATCH 1/7] feat: implement distributed gradient norm clipping --- docs/gradient_norm_clipping.md | 22 +++ .../nn/parallel/ddp/distributed_optimizer.h | 9 ++ .../nn/parallel/ddp/param_and_grad_buffer.h | 2 + infini_train/include/optimizer.h | 27 ++++ .../src/kernels/cpu/accumulate_grad.cc | 29 ++++ .../src/kernels/cuda/accumulate_grad.cu | 26 ++++ .../nn/parallel/ddp/distributed_optimizer.cc | 122 ++++++++++++++- infini_train/src/optimizer.cc | 138 ++++++++++++++++ tests/optimizer/CMakeLists.txt | 52 +++++++ tests/optimizer/test_clip_grad_norm.cc | 118 ++++++++++++++ .../test_optimizer_parameter_names.cc | 147 ++++++++++++++++++ 11 files changed, 689 insertions(+), 3 deletions(-) create mode 100644 docs/gradient_norm_clipping.md create mode 100644 tests/optimizer/test_clip_grad_norm.cc diff --git a/docs/gradient_norm_clipping.md b/docs/gradient_norm_clipping.md new file mode 100644 index 000000000..9966517f5 --- /dev/null +++ b/docs/gradient_norm_clipping.md @@ -0,0 +1,22 @@ +# Distributed gradient norm and clipping + +The public entry point is: + + auto total_norm = optimizer->ClipGradNorm_(parameters, max_norm, norm_type, error_if_nonfinite, std::nullopt); + +It returns a CPU FP32 scalar containing the pre-clipping norm and scales the selected gradients in place. Parameters without gradients are ignored and duplicate parameter pointers are counted once. Supported norm types are positive finite p and positive infinity. The coefficient is min(max_norm / (total_norm + 1e-6), 1); max_norm == 0 therefore zeros gradients rather than disabling clipping. + +foreach=true is intentionally rejected because a fused multi-tensor kernel is not part of this implementation yet. The CPU and CUDA ScaleInplace kernels write through the existing gradient view, preserving flat-buffer and ZeRO-shard aliases. CUDA norm inspection synchronizes only to obtain scalar statistics; scaling remains on the CUDA stream. + +DistributedOptimizer::ClipGradNorm_ first finishes pending gradient collectives, computes the norm over the shard parameters consumed by its base optimizer, reduces statistics over the DP process group, and applies the same coefficient to every local shard. When PP is enabled, statistics are also reduced over the PP process group. Pipeline scheduling calls ClipGradNormConfigured after all micro-batches complete. + +The GPT-2 and LLaMA-3 examples expose: + +- --clip_grad_norm=-1 (negative disables clipping) +- --grad_norm_type=2 +- --clip_grad_error_if_nonfinite=true +- --clip_grad_foreach=auto|true|false + +A total_grad_norm log line is emitted when clipping is enabled. + +Current limitation: tensor-parallel replicated-parameter ownership metadata is not yet represented by Tensor/Optimizer, so TP/SP duplicate filtering needs follow-up integration before claiming full mixed-parallel equivalence. diff --git a/infini_train/include/nn/parallel/ddp/distributed_optimizer.h b/infini_train/include/nn/parallel/ddp/distributed_optimizer.h index d7cea198e..7edacb3eb 100644 --- a/infini_train/include/nn/parallel/ddp/distributed_optimizer.h +++ b/infini_train/include/nn/parallel/ddp/distributed_optimizer.h @@ -3,7 +3,9 @@ #include #include #include +#include #include +#include #include #include "infini_train/include/optimizer.h" @@ -31,6 +33,10 @@ class DistributedOptimizer final : public infini_train::Optimizer { void Step() override; + std::shared_ptr ClipGradNorm_( + const std::vector> ¶meters, float max_norm, float norm_type = 2.0f, + bool error_if_nonfinite = false, std::optional foreach = std::nullopt) override; + void ZeroGrad(bool set_to_none = true) override; std::unordered_map> StateDict() const override; @@ -63,6 +69,9 @@ class DistributedOptimizer final : public infini_train::Optimizer { // Base optimizer (SGD, Adam and etc.) std::shared_ptr base_optimizer_; + std::vector> shard_params_; + std::vector> shard_param_owners_; + std::unordered_set tp_replicated_param_owners_; }; } // namespace infini_train::nn::parallel diff --git a/infini_train/include/nn/parallel/ddp/param_and_grad_buffer.h b/infini_train/include/nn/parallel/ddp/param_and_grad_buffer.h index 2c572984d..e29b08467 100644 --- a/infini_train/include/nn/parallel/ddp/param_and_grad_buffer.h +++ b/infini_train/include/nn/parallel/ddp/param_and_grad_buffer.h @@ -125,6 +125,8 @@ class ParamAndGradBucketGroup { const DistributedDataParallelConfig &config() const { return ddp_config_; } + const ProcessGroup *collective_pg() const { return collective_pg_; } + private: std::vector> buckets_; const ProcessGroup *collective_pg_ = nullptr; diff --git a/infini_train/include/optimizer.h b/infini_train/include/optimizer.h index d85b1acea..90d49f430 100644 --- a/infini_train/include/optimizer.h +++ b/infini_train/include/optimizer.h @@ -3,6 +3,7 @@ #include #include #include +#include #include #include #include @@ -27,6 +28,25 @@ class Optimizer { virtual void ZeroGrad(bool set_to_none = true); + // Return the pre-clipping norm and scale selected gradients in place. + virtual std::shared_ptr ClipGradNorm_( + const std::vector> ¶meters, float max_norm, float norm_type = 2.0f, + bool error_if_nonfinite = false, std::optional foreach = std::nullopt); + + std::shared_ptr ClipGradNorm(const std::vector> ¶meters, float max_norm, + float norm_type = 2.0f, bool error_if_nonfinite = false, + std::optional foreach = std::nullopt) { + return ClipGradNorm_(parameters, max_norm, norm_type, error_if_nonfinite, foreach); + } + + // Scale selected gradients without replacing their storage. + void ScaleGradients_(const std::vector> ¶meters, float scale); + + void SetClipGradNormConfig(float max_norm, float norm_type = 2.0f, bool error_if_nonfinite = false, + std::optional foreach = std::nullopt); + bool HasClipGradNormConfig() const { return clip_grad_norm_config_.has_value(); } + std::shared_ptr ClipGradNormConfigured(); + virtual void Step() = 0; virtual std::unordered_map> StateDict() const { return {}; }; @@ -49,6 +69,13 @@ class Optimizer { float learning_rate_ = 0.0f; float initial_learning_rate_ = 0.0f; bool initial_lr_set_ = false; + struct ClipGradNormConfig { + float max_norm; + float norm_type; + bool error_if_nonfinite; + std::optional foreach; + }; + std::optional clip_grad_norm_config_; }; namespace optimizers { diff --git a/infini_train/src/kernels/cpu/accumulate_grad.cc b/infini_train/src/kernels/cpu/accumulate_grad.cc index cfe85b9c2..a78b5175b 100644 --- a/infini_train/src/kernels/cpu/accumulate_grad.cc +++ b/infini_train/src/kernels/cpu/accumulate_grad.cc @@ -5,6 +5,34 @@ #include "infini_train/include/tensor.h" namespace infini_train::kernels::cpu { + +template +void ScaleInplaceTyped(const std::shared_ptr &tensor, float scale) { + auto *data = static_cast(tensor->DataPtr()); +#pragma omp parallel for + for (size_t i = 0; i < tensor->NumElements(); ++i) { + data[i] = T(static_cast(data[i]) * scale); + } +} + +void ScaleInplace(const std::shared_ptr &tensor, float scale) { + switch (tensor->Dtype()) { + case DataType::kFLOAT16: + ScaleInplaceTyped(tensor, scale); + return; + case DataType::kBFLOAT16: + ScaleInplaceTyped(tensor, scale); + return; + case DataType::kFLOAT32: + ScaleInplaceTyped(tensor, scale); + return; + case DataType::kFLOAT64: + ScaleInplaceTyped(tensor, scale); + return; + default: + LOG(FATAL) << "ScaleInplace only supports floating point gradients."; + } +} void AccumulateGrad(const std::shared_ptr &gradient, float rate, const std::shared_ptr &tensor) { for (int64_t idx = 0; idx < gradient->NumElements(); ++idx) { static_cast(tensor->DataPtr())[idx] += rate * static_cast(gradient->DataPtr())[idx]; @@ -40,6 +68,7 @@ void AdamAccumulateGrad(const std::shared_ptr &grad, const std::shared_p REGISTER_KERNEL(infini_train::Device::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) REGISTER_CPU_ACCUMULATE_GRAD_KERNEL(AccumulateGrad) +REGISTER_CPU_ACCUMULATE_GRAD_KERNEL(ScaleInplace) REGISTER_CPU_ACCUMULATE_GRAD_KERNEL(AdamAccumulateGrad) #undef REGISTER_CPU_ACCUMULATE_GRAD_KERNEL diff --git a/infini_train/src/kernels/cuda/accumulate_grad.cu b/infini_train/src/kernels/cuda/accumulate_grad.cu index 93409a7ef..b107b37ef 100644 --- a/infini_train/src/kernels/cuda/accumulate_grad.cu +++ b/infini_train/src/kernels/cuda/accumulate_grad.cu @@ -11,6 +11,31 @@ namespace infini_train::kernels::cuda { +template +__global__ void ScaleInplaceKernel(T *data, float scale, size_t num_elements) { + const size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < num_elements) { + data[idx] = common::cuda::Mul(data[idx], common::cuda::Cast(scale)); + } +} + +void ScaleInplace(const std::shared_ptr &tensor, float scale) { + const size_t num_elements = tensor->NumElements(); + const int threads_per_block = 256; + const int num_blocks = (num_elements + threads_per_block - 1) / threads_per_block; + auto device = tensor->GetDevice(); + const auto &cuda_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); + core::cuda::DispatchCudaFunc( + tensor->Dtype(), + [=]() { + ScaleInplaceKernel<<>>( + static_cast(tensor->DataPtr()), scale, num_elements); + }, + "CUDA ScaleInplace"); +} + template __global__ void AccumulateGradKernel(const T *grad_ptr, float rate, T *tensor_ptr, size_t num_elements) { int idx = blockIdx.x * blockDim.x + threadIdx.x; @@ -90,6 +115,7 @@ void AdamAccumulateGrad(const std::shared_ptr &grad, const std::shared_p REGISTER_KERNEL(infini_train::Device::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL(AccumulateGrad) +REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL(ScaleInplace) REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL(AdamAccumulateGrad) #undef REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL diff --git a/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc b/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc index 523bcf2d7..afb10649c 100644 --- a/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc +++ b/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc @@ -1,8 +1,15 @@ #include "infini_train/include/nn/parallel/ddp/distributed_optimizer.h" +#include +#include +#include + #include "glog/logging.h" #include "infini_train/include/nn/parallel/ddp/distributed_data_parallel.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/nn/parallel/tensor_parallel.h" +#include "infini_train/include/nn/parallel/utils.h" #include "infini_train/include/tensor.h" namespace infini_train::nn::parallel { @@ -38,12 +45,14 @@ DistributedOptimizer::DistributedOptimizer(OptimizerCreatorNamed creator, const NamedParameterList shard_named_parameters; BuildShardParamsAndBindGrads( - [¶meter_name_by_tensor, &shard_named_parameters](const std::shared_ptr ¶meter, + [this, ¶meter_name_by_tensor, &shard_named_parameters](const std::shared_ptr ¶meter, const std::shared_ptr ¶m_piece) { const auto name_it = parameter_name_by_tensor.find(parameter.get()); CHECK(name_it != parameter_name_by_tensor.end()) << "DistributedOptimizer parameter is not registered in the model"; shard_named_parameters.emplace_back(name_it->second, param_piece); + shard_params_.push_back(param_piece); + shard_param_owners_.push_back(parameter); }); base_optimizer_ = creator(shard_named_parameters); @@ -61,6 +70,15 @@ void DistributedOptimizer::InitializeModelChunks(const std::vectorparam_grad_buffers().end()); bucket_groups_.insert(bucket_groups_.end(), ddp_chunk->bucket_groups().begin(), ddp_chunk->bucket_groups().end()); + if (global::GetTensorParallelSize() > 1) { + for (const auto &[_, module] : ddp_chunk->module()->NamedModules()) { + auto row_parallel = std::dynamic_pointer_cast(module); + if (row_parallel && row_parallel->bias()) { + tp_replicated_param_owners_.insert( + row_parallel->parameter(RowParallelLinear::kParamBiasName).get()); + } + } + } } } @@ -121,8 +139,10 @@ void DistributedOptimizer::BuildShardParamsAndBindGrads(const AddShardParam &add } } - CHECK_GT(num_shard_params, 0) << "DistributedOptimizer: this DP rank owns no param pieces. " - << "Check bucket padding/divisibility and param bucketing order."; + // A rank may legitimately own no non-padding elements when a parameter is + // smaller than the DP shard size. Keep an empty base optimizer on that + // rank; global norm reduction still participates with a zero local stat. + (void)num_shard_params; } void DistributedOptimizer::StartGradSync() { @@ -141,6 +161,102 @@ void DistributedOptimizer::FinishParamSync(bool skip_next_bucket_dispatch) { for (auto &group : bucket_groups_) { group->FinishParamSync(skip_next_bucket_dispatch); } } +std::shared_ptr DistributedOptimizer::ClipGradNorm_( + const std::vector> ¶meters, float max_norm, float norm_type, + bool error_if_nonfinite, std::optional foreach) { + CHECK_GE(max_norm, 0.0f) << "max_norm must be non-negative."; + CHECK((norm_type > 0.0f && std::isfinite(norm_type)) || norm_type == std::numeric_limits::infinity()) + << "norm_type must be positive finite or +inf."; + if (foreach.value_or(false)) { + LOG(FATAL) << "DistributedOptimizer: foreach=true is not implemented."; + } + + FinishGradSync(); + + std::unordered_set requested; + for (const auto ¶meter : parameters) { + if (parameter) { + requested.insert(parameter.get()); + } + } + std::vector> selected_shards; + std::vector> norm_shards; + for (size_t i = 0; i < shard_params_.size(); ++i) { + if (requested.empty() || requested.contains(shard_param_owners_[i].get())) { + selected_shards.push_back(shard_params_[i]); + bool include_in_norm = true; + if (global::GetTensorParallelSize() > 1 && + tp_replicated_param_owners_.contains(shard_param_owners_[i].get())) { + const auto *tp_group = ProcessGroupFactory::Instance(shard_params_[i]->GetDevice().type()) + ->Get(GetTensorParallelProcessGroupName( + shard_params_[i]->GetDevice().Rank().GlobalRank())); + CHECK(tp_group) << "Tensor-parallel process group is not initialized."; + include_in_norm = tp_group->GetGroupRank(shard_params_[i]->GetDevice().Rank().GlobalRank()) == 0; + } + if (include_in_norm) { + norm_shards.push_back(shard_params_[i]); + } + } + } + + // Ask the base optimizer for the local norm without scaling. Infinite + // max_norm makes its coefficient exactly one for finite gradients. + auto local_norm_tensor = base_optimizer_->ClipGradNorm_( + norm_shards, std::numeric_limits::infinity(), norm_type, false, std::nullopt); + const float local_norm = *static_cast(local_norm_tensor->DataPtr()); + const bool is_inf_norm = std::isinf(norm_type); + double local_stat = is_inf_norm ? static_cast(local_norm) + : std::pow(static_cast(local_norm), static_cast(norm_type)); + + const ProcessGroup *group = nullptr; + for (const auto &bucket_group : bucket_groups_) { + if (bucket_group->collective_pg()) { + group = bucket_group->collective_pg(); + break; + } + } + Device total_norm_device = selected_shards.empty() ? Device() : selected_shards.front()->GetDevice(); + if (selected_shards.empty()) { + for (const auto &bucket_group : bucket_groups_) { + if (!bucket_group->buckets().empty() && bucket_group->buckets().front()->param_data()) { + total_norm_device = bucket_group->buckets().front()->param_data()->GetDevice(); + break; + } + } + } + auto reduced = std::make_shared(std::vector{}, DataType::kFLOAT32, total_norm_device); + reduced->Fill(static_cast(local_stat)); + if (group && ddp_world_size_ > 1) { + group->AllReduce(reduced, is_inf_norm ? function::ReduceOpType::kMax : function::ReduceOpType::kSum, false); + } + if (global::GetTensorParallelSize() > 1) { + const auto *tp_group = ProcessGroupFactory::Instance(total_norm_device.type())->Get( + GetTensorParallelProcessGroupName(total_norm_device.Rank().GlobalRank())); + CHECK(tp_group) << "Tensor-parallel process group is not initialized."; + tp_group->AllReduce(reduced, is_inf_norm ? function::ReduceOpType::kMax : function::ReduceOpType::kSum, false); + } + if (nn::parallel::global::GetPipelineParallelSize() > 1) { + const auto *pp_group = ProcessGroupFactory::Instance(total_norm_device.type())->Get( + GetPipelineParallelProcessGroupName(total_norm_device.Rank().GlobalRank())); + CHECK(pp_group) << "Pipeline process group is not initialized."; + pp_group->AllReduce(reduced, is_inf_norm ? function::ReduceOpType::kMax : function::ReduceOpType::kSum, false); + } + Tensor reduced_cpu = reduced->GetDevice().IsCPU() ? Tensor(*reduced, 0, reduced->Dims()) : reduced->To(Device()); + const float reduced_value = *static_cast(reduced_cpu.DataPtr()); + const double total_norm = is_inf_norm ? static_cast(reduced_value) + : std::pow(static_cast(reduced_value), 1.0 / norm_type); + if (error_if_nonfinite && !std::isfinite(total_norm)) { + LOG(FATAL) << "The total gradient norm is non-finite."; + } + + const double coefficient = std::min(static_cast(max_norm) / (total_norm + 1e-6), 1.0); + base_optimizer_->ScaleGradients_(selected_shards, static_cast(coefficient)); + + auto result = std::make_shared(std::vector{}, DataType::kFLOAT32, Device()); + *static_cast(result->DataPtr()) = static_cast(total_norm); + return result; +} + void DistributedOptimizer::ZeroGrad(bool set_to_none) { // Clear BucketGroup state and reset buffer: // If set_to_none is true: diff --git a/infini_train/src/optimizer.cc b/infini_train/src/optimizer.cc index 39b999c77..f57219537 100644 --- a/infini_train/src/optimizer.cc +++ b/infini_train/src/optimizer.cc @@ -1,11 +1,18 @@ #include "infini_train/include/optimizer.h" +#include +#include +#include #include +#include #include #include "infini_train/include/core/runtime/device_guard.h" #include "infini_train/include/device.h" #include "infini_train/include/dispatcher.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/nn/parallel/process_group.h" +#include "infini_train/include/nn/parallel/utils.h" #include "infini_train/include/tensor.h" namespace infini_train { @@ -29,6 +36,137 @@ void Optimizer::ZeroGrad(bool set_to_none) { for (auto param : params_) { param->ZeroGrad(set_to_none); } } +namespace { +struct GradientNormStats { + double sum = 0.0; + double max_abs = 0.0; + bool has_gradient = false; + bool finite = true; +}; + +GradientNormStats ComputeGradientNormStats(const std::vector> ¶meters, float norm_type) { + GradientNormStats stats; + std::unordered_set seen; + for (const auto ¶meter : parameters) { + if (!parameter || !seen.insert(parameter.get()).second || !parameter->grad()) { + continue; + } + stats.has_gradient = true; + const auto &gradient = parameter->grad(); + Tensor host = gradient->GetDevice().IsCPU() ? Tensor(*gradient, 0, gradient->Dims()) : gradient->To(Device()); + if (gradient->GetDevice().IsCUDA()) { + core::GetDeviceGuardImpl(gradient->GetDevice().type())->SynchronizeDevice(gradient->GetDevice()); + } + Tensor fp32 = host.Dtype() == DataType::kFLOAT32 ? host : host.To(DataType::kFLOAT32); + const float *data = static_cast(fp32.DataPtr()); + for (size_t i = 0; i < fp32.NumElements(); ++i) { + const double value = static_cast(data[i]); + const double abs_value = std::abs(value); + if (!std::isfinite(value)) { + stats.finite = false; + if (std::isinf(norm_type)) { + stats.max_abs = std::numeric_limits::quiet_NaN(); + } + } + if (std::isinf(norm_type)) { + stats.max_abs = std::max(stats.max_abs, abs_value); + } else { + stats.sum += std::pow(abs_value, static_cast(norm_type)); + } + } + } + return stats; +} + +std::shared_ptr MakeScalar(float value) { + auto result = std::make_shared(std::vector{}, DataType::kFLOAT32, Device()); + *static_cast(result->DataPtr()) = value; + return result; +} + +void ScaleGradientInplace(const std::shared_ptr &gradient, float scale) { + if (!gradient || scale == 1.0f) { + return; + } + auto device = gradient->GetDevice(); + core::DeviceGuard guard(device); + auto kernel = Dispatcher::Instance().GetKernel({device.type(), "ScaleInplace"}); + kernel.Call(gradient, scale); +} +} // namespace + +void Optimizer::SetClipGradNormConfig(float max_norm, float norm_type, bool error_if_nonfinite, + std::optional foreach) { + clip_grad_norm_config_ = ClipGradNormConfig{max_norm, norm_type, error_if_nonfinite, foreach}; +} + +std::shared_ptr Optimizer::ClipGradNormConfigured() { + if (!clip_grad_norm_config_) { + return nullptr; + } + const auto &config = *clip_grad_norm_config_; + return ClipGradNorm_(params_, config.max_norm, config.norm_type, config.error_if_nonfinite, config.foreach); +} + +void Optimizer::ScaleGradients_(const std::vector> ¶meters, float scale) { + std::unordered_set seen; + for (const auto ¶meter : parameters) { + if (!parameter || !seen.insert(parameter.get()).second || !parameter->grad()) { + continue; + } + ScaleGradientInplace(parameter->grad(), scale); + } +} + +std::shared_ptr Optimizer::ClipGradNorm_(const std::vector> ¶meters, + float max_norm, float norm_type, bool error_if_nonfinite, + std::optional foreach) { + CHECK_GE(max_norm, 0.0f) << "max_norm must be non-negative."; + CHECK((norm_type > 0.0f && std::isfinite(norm_type)) || norm_type == std::numeric_limits::infinity()) + << "norm_type must be positive finite or +inf."; + if (foreach.value_or(false)) { + LOG(FATAL) << "ClipGradNorm_: foreach=true is not implemented."; + } + + const auto stats = ComputeGradientNormStats(parameters, norm_type); + double total_norm = 0.0; + if (stats.has_gradient) { + total_norm = std::isinf(norm_type) ? stats.max_abs : std::pow(stats.sum, 1.0 / norm_type); + } + + // Pipeline stages own disjoint parameter sets. Reduce their statistics so + // every stage uses one global clipping coefficient. + if (stats.has_gradient && infini_train::nn::parallel::global::GetPipelineParallelSize() > 1) { + const auto device = [&]() { + for (const auto ¶meter : parameters) { + if (parameter && parameter->grad()) return parameter->grad()->GetDevice(); + } + return Device(); + }(); + auto reduced = std::make_shared(std::vector{}, DataType::kFLOAT32, device); + reduced->Fill(static_cast(std::isinf(norm_type) ? stats.max_abs : stats.sum)); + const auto *group = nn::parallel::ProcessGroupFactory::Instance(device.type())->Get( + nn::parallel::GetPipelineParallelProcessGroupName(device.Rank().GlobalRank())); + CHECK(group) << "Pipeline process group is not initialized."; + group->AllReduce(reduced, std::isinf(norm_type) ? nn::parallel::function::ReduceOpType::kMax + : nn::parallel::function::ReduceOpType::kSum, false); + Tensor reduced_cpu = reduced->GetDevice().IsCPU() ? Tensor(*reduced, 0, reduced->Dims()) : reduced->To(Device()); + const float value = *static_cast(reduced_cpu.DataPtr()); + total_norm = std::isinf(norm_type) ? value : std::pow(static_cast(value), 1.0 / norm_type); + } + if (error_if_nonfinite && (!std::isfinite(total_norm))) { + LOG(FATAL) << "The total gradient norm is non-finite."; + } + + const double coefficient = std::isinf(max_norm) + ? 1.0 + : std::min(static_cast(max_norm) / (total_norm + 1e-6), 1.0); + if (stats.has_gradient) { + ScaleGradients_(parameters, static_cast(coefficient)); + } + return MakeScalar(static_cast(total_norm)); +} + void Optimizer::set_learning_rate(float lr) { learning_rate_ = lr; } float Optimizer::learning_rate() const { return learning_rate_; } diff --git a/tests/optimizer/CMakeLists.txt b/tests/optimizer/CMakeLists.txt index c0bfbd503..a562b0146 100644 --- a/tests/optimizer/CMakeLists.txt +++ b/tests/optimizer/CMakeLists.txt @@ -7,3 +7,55 @@ file(GLOB OPTIMIZER_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/test_*.cc) infini_train_add_test_suite(test_optimizer SOURCES ${OPTIMIZER_SOURCES} ) + +add_test( + NAME OptimizerParameterNamesTest.DistributedOptimizerPropagatesNamesToShardOptimizer + COMMAND ${CMAKE_COMMAND} -E env INFINI_RUN_ID=optimizer_ctest torchrun --standalone --no-python --nproc_per_node=2 + $ + --gtest_filter=CUDA/OptimizerParameterNamesTest.DistributedOptimizerPropagatesNamesToShardOptimizer/* +) +set_tests_properties( + OptimizerParameterNamesTest.DistributedOptimizerPropagatesNamesToShardOptimizer + PROPERTIES + LABELS "cuda;distributed" + TIMEOUT 30 +) + +add_test( + NAME OptimizerParameterNamesTest.DistributedOptimizerClipGradNormUsesGlobalShardNorm + COMMAND ${CMAKE_COMMAND} -E env INFINI_RUN_ID=optimizer_ctest torchrun --standalone --no-python --nproc_per_node=2 + $ + --gtest_filter=CUDA/OptimizerParameterNamesTest.DistributedOptimizerClipGradNormUsesGlobalShardNorm/* +) +set_tests_properties( + OptimizerParameterNamesTest.DistributedOptimizerClipGradNormUsesGlobalShardNorm + PROPERTIES + LABELS "cuda;distributed" + TIMEOUT 30 +) + +add_test( + NAME OptimizerParameterNamesTest.DistributedOptimizerClipGradNormUsesZero2LocalShard + COMMAND ${CMAKE_COMMAND} -E env INFINI_RUN_ID=optimizer_zero2_ctest torchrun --standalone --no-python --nproc_per_node=2 + $ + --gtest_filter=CUDA/OptimizerParameterNamesTest.DistributedOptimizerClipGradNormUsesZero2LocalShard/* +) +set_tests_properties( + OptimizerParameterNamesTest.DistributedOptimizerClipGradNormUsesZero2LocalShard + PROPERTIES + LABELS "cuda;distributed" + TIMEOUT 30 +) + +add_test( + NAME OptimizerParameterNamesTest.DistributedOptimizerClipGradNormHandlesEmptyLocalShard + COMMAND ${CMAKE_COMMAND} -E env INFINI_RUN_ID=optimizer_empty_ctest torchrun --standalone --no-python --nproc_per_node=2 + $ + --gtest_filter=CUDA/OptimizerParameterNamesTest.DistributedOptimizerClipGradNormHandlesEmptyLocalShard/* +) +set_tests_properties( + OptimizerParameterNamesTest.DistributedOptimizerClipGradNormHandlesEmptyLocalShard + PROPERTIES + LABELS "cuda;distributed" + TIMEOUT 30 +) diff --git a/tests/optimizer/test_clip_grad_norm.cc b/tests/optimizer/test_clip_grad_norm.cc new file mode 100644 index 000000000..6a690258d --- /dev/null +++ b/tests/optimizer/test_clip_grad_norm.cc @@ -0,0 +1,118 @@ +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/optimizer.h" +#include "infini_train/include/tensor.h" +#include "tests/common/test_utils.h" + +using namespace infini_train; + +class ClipGradNormTest : public infini_train::test::InfiniTrainTest {}; + +static std::shared_ptr MakeTensor(Device device, const std::vector &values) { + auto tensor = std::make_shared(std::vector{static_cast(values.size())}, + DataType::kFLOAT32, device); + if (device.IsCPU()) { + std::copy(values.begin(), values.end(), static_cast(tensor->DataPtr())); + } else { + auto cpu = std::make_shared(std::vector{static_cast(values.size())}, + DataType::kFLOAT32, Device()); + std::copy(values.begin(), values.end(), static_cast(cpu->DataPtr())); + tensor->CopyFrom(*cpu); + } + return tensor; +} + +static float ScalarCPU(const std::shared_ptr &value) { + auto cpu = value->To(Device()); + return *static_cast(cpu.DataPtr()); +} + +TEST_P(ClipGradNormTest, L2ClipsInPlace) { + auto param = std::make_shared(std::vector{2}, DataType::kFLOAT32, GetDevice()); + auto grad = MakeTensor(GetDevice(), {3.0f, 4.0f}); + param->set_grad(grad); + auto optimizer = std::make_shared( + std::vector>{param}, 0.1f); + + auto total_norm = optimizer->ClipGradNorm({param}, 2.0f, 2.0f); + EXPECT_NEAR(ScalarCPU(total_norm), 5.0f, 1e-5f); + auto clipped = grad->To(Device()); + const float *values = static_cast(clipped.DataPtr()); + EXPECT_NEAR(values[0], 1.2f, 1e-5f); + EXPECT_NEAR(values[1], 1.6f, 1e-5f); +} + +TEST_P(ClipGradNormTest, SupportsL1AndInfinity) { + auto param = std::make_shared(std::vector{3}, DataType::kFLOAT32, GetDevice()); + auto grad = MakeTensor(GetDevice(), {-2.0f, 1.0f, 3.0f}); + param->set_grad(grad); + auto optimizer = std::make_shared( + std::vector>{param}, 0.1f); + + auto l1 = optimizer->ClipGradNorm({param}, 3.0f, 1.0f); + EXPECT_NEAR(ScalarCPU(l1), 6.0f, 1e-5f); + auto after_l1 = grad->To(Device()); + const float *l1_values = static_cast(after_l1.DataPtr()); + EXPECT_NEAR(l1_values[0], -1.0f, 1e-5f); + EXPECT_NEAR(l1_values[1], 0.5f, 1e-5f); + EXPECT_NEAR(l1_values[2], 1.5f, 1e-5f); + + auto inf_grad = MakeTensor(GetDevice(), {-2.0f, 1.0f, 3.0f}); + param->set_grad(inf_grad); + auto inf = optimizer->ClipGradNorm({param}, 1.0f, std::numeric_limits::infinity()); + EXPECT_NEAR(ScalarCPU(inf), 3.0f, 1e-5f); + auto after_inf = inf_grad->To(Device()); + const float *inf_values = static_cast(after_inf.DataPtr()); + EXPECT_NEAR(inf_values[0], -2.0f / 3.0f, 1e-5f); + EXPECT_NEAR(inf_values[2], 1.0f, 1e-5f); +} + +TEST_P(ClipGradNormTest, IgnoresMissingAndDuplicateGradients) { + auto first = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice()); + first->set_grad(MakeTensor(GetDevice(), {3.0f})); + auto second = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice()); + auto optimizer = std::make_shared( + std::vector>{first, second}, 0.1f); + + auto total_norm = optimizer->ClipGradNorm({first, first, second}, 1.0f, 2.0f); + EXPECT_NEAR(ScalarCPU(total_norm), 3.0f, 1e-5f); + auto clipped = first->grad()->To(Device()); + EXPECT_NEAR(*static_cast(clipped.DataPtr()), 1.0f, 1e-5f); + + auto empty = optimizer->ClipGradNorm({second}, 1.0f, 2.0f); + EXPECT_FLOAT_EQ(ScalarCPU(empty), 0.0f); +} + +INFINI_TRAIN_REGISTER_TEST(ClipGradNormTest); + +TEST_P(ClipGradNormTest, NonIntegerPAndZeroMaxNorm) { + auto param = std::make_shared(std::vector{2}, DataType::kFLOAT32, GetDevice()); + auto grad = MakeTensor(GetDevice(), {2.0f, 2.0f}); + param->set_grad(grad); + auto optimizer = std::make_shared( + std::vector>{param}, 0.1f); + + const float expected = std::pow(2.0f * std::pow(2.0f, 3.5f), 1.0f / 3.5f); + auto total_norm = optimizer->ClipGradNorm({param}, 0.0f, 3.5f); + EXPECT_NEAR(ScalarCPU(total_norm), expected, 1e-5f); + auto clipped = grad->To(Device()); + const float *values = static_cast(clipped.DataPtr()); + EXPECT_FLOAT_EQ(values[0], 0.0f); + EXPECT_FLOAT_EQ(values[1], 0.0f); +} + +TEST_P(ClipGradNormTest, ErrorOnNonFiniteBeforeScaling) { + ONLY_CPU(); + auto param = std::make_shared(std::vector{2}, DataType::kFLOAT32, GetDevice()); + auto grad = MakeTensor(GetDevice(), {std::numeric_limits::quiet_NaN(), 1.0f}); + param->set_grad(grad); + auto optimizer = std::make_shared( + std::vector>{param}, 0.1f); + EXPECT_DEATH(optimizer->ClipGradNorm({param}, 1.0f, 2.0f, true), "non-finite"); +} diff --git a/tests/optimizer/test_optimizer_parameter_names.cc b/tests/optimizer/test_optimizer_parameter_names.cc index 3b50b5858..f26ede2f2 100644 --- a/tests/optimizer/test_optimizer_parameter_names.cc +++ b/tests/optimizer/test_optimizer_parameter_names.cc @@ -3,6 +3,14 @@ #include "gtest/gtest.h" +#include "infini_train/include/nn/modules/linear.h" +#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel.h" +#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel_config.h" +#include "infini_train/include/nn/parallel/ddp/distributed_optimizer.h" +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/nn/parallel/process_group.h" +#include "infini_train/include/nn/parallel/rank.h" +#include "infini_train/include/nn/parallel/utils.h" #include "infini_train/include/optimizer.h" #include "infini_train/include/tensor.h" @@ -44,6 +52,42 @@ TEST_P(OptimizerParameterNamesTest, ConstructorMatchesNamesToOptimizerParameterO EXPECT_TRUE(state.contains("adam.v.first")); } +TEST_P(OptimizerParameterNamesTest, DistributedOptimizerPropagatesNamesToShardOptimizer) { + ONLY_CUDA(); + REQUIRE_MIN_DEVICES(2); + if (nn::parallel::global::GetDataParallelSize() != 2) { + GTEST_SKIP() << "requires PROC_WORLD_SIZE=2"; + } + + const nn::parallel::Rank rank(nn::parallel::global::GetGlobalProcRank(), /*thread_rank=*/0, + nn::parallel::global::GetNprocPerNode(), /*thread_size=*/1); + auto *pg_factory = nn::parallel::ProcessGroupFactory::Instance(Device::DeviceType::kCUDA); + pg_factory->GetOrCreate(nn::parallel::GetDataParallelProcessGroupName(rank.GlobalRank()), + nn::parallel::GetDataParallelGroupRanks(rank.GlobalRank())); + + auto model = std::make_shared(64, 4, /*bias=*/false, + Device(Device::DeviceType::kCUDA, + nn::parallel::global::GetLocalProcRank())); + const auto named_parameters = model->NamedParameters(); + + nn::parallel::DistributedDataParallelConfig ddp_config; + ddp_config.zero_stage = 1; + ddp_config.overlap_grad_reduce = false; + ddp_config.overlap_param_gather = false; + auto ddp_model = std::make_shared(model, rank, ddp_config); + + nn::parallel::DistributedOptimizer optimizer(optimizers::Adam::CreateNamed(0.001), named_parameters, + std::vector>{ddp_model}, + /*ddp_world_size=*/2, + /*ddp_rank=*/rank.GlobalRank()); + const auto state = optimizer.StateDict(); + + EXPECT_TRUE(state.contains("adam.m.weight")); + EXPECT_TRUE(state.contains("adam.v.weight")); + EXPECT_FALSE(state.contains("adam.m.0")); + EXPECT_FALSE(state.contains("adam.v.0")); +} + TEST_P(OptimizerParameterNamesTest, PreservesNumericKeysWhenNamesAreNotSet) { auto parameter = std::make_shared(std::vector{2, 2}, DataType::kFLOAT32, GetDevice()); auto adam = std::make_shared(std::vector>{parameter}, 0.001); @@ -53,4 +97,107 @@ TEST_P(OptimizerParameterNamesTest, PreservesNumericKeysWhenNamesAreNotSet) { EXPECT_TRUE(state.contains("adam.v.0")); } + +TEST_P(OptimizerParameterNamesTest, DistributedOptimizerClipGradNormUsesZero2LocalShard) { + ONLY_CUDA(); + REQUIRE_MIN_DEVICES(2); + if (nn::parallel::global::GetDataParallelSize() != 2) { + GTEST_SKIP() << "requires WORLD_SIZE=2"; + } + const nn::parallel::Rank rank(nn::parallel::global::GetGlobalProcRank(), /*thread_rank=*/0, + nn::parallel::global::GetNprocPerNode(), /*thread_size=*/1); + auto *pg_factory = nn::parallel::ProcessGroupFactory::Instance(Device::DeviceType::kCUDA); + pg_factory->GetOrCreate(nn::parallel::GetDataParallelProcessGroupName(rank.GlobalRank()), + nn::parallel::GetDataParallelGroupRanks(rank.GlobalRank())); + auto model = std::make_shared(64, 4, /*bias=*/false, + Device(Device::DeviceType::kCUDA, + nn::parallel::global::GetLocalProcRank())); + nn::parallel::DistributedDataParallelConfig ddp_config; + ddp_config.zero_stage = 2; + ddp_config.overlap_grad_reduce = false; + ddp_config.overlap_param_gather = false; + auto ddp_model = std::make_shared(model, rank, ddp_config); + const auto named_parameters = model->NamedParameters(); + nn::parallel::DistributedOptimizer optimizer(optimizers::SGD::CreateNamed(0.1f), named_parameters, + std::vector>{ddp_model}, + /*ddp_world_size=*/2, /*ddp_rank=*/rank.GlobalRank()); + auto parameter = model->Parameters().front(); + auto full_grad = std::make_shared(parameter->Dims(), DataType::kFLOAT32, parameter->GetDevice()); + full_grad->Fill(1.0f); + auto group = ddp_model->bucket_groups().front(); + group->AccumulateParamGrad(parameter, full_grad, /*overwrite=*/true, /*learning_rate=*/1.0f); + auto total_norm = optimizer.ClipGradNorm(model->Parameters(), 8.0f, 2.0f); + auto total_norm_cpu = total_norm->To(Device()); + EXPECT_NEAR(*static_cast(total_norm_cpu.DataPtr()), 16.0f, 1e-4f); + auto local_shard = group->GetLocalGradShardBuffer(0)->To(Device()); + EXPECT_NEAR(static_cast(local_shard.DataPtr())[0], 0.5f, 1e-4f); +} + + +TEST_P(OptimizerParameterNamesTest, DistributedOptimizerClipGradNormHandlesEmptyLocalShard) { + ONLY_CUDA(); + REQUIRE_MIN_DEVICES(2); + if (nn::parallel::global::GetDataParallelSize() != 2) { + GTEST_SKIP() << "requires WORLD_SIZE=2"; + } + const nn::parallel::Rank rank(nn::parallel::global::GetGlobalProcRank(), /*thread_rank=*/0, + nn::parallel::global::GetNprocPerNode(), /*thread_size=*/1); + auto *pg_factory = nn::parallel::ProcessGroupFactory::Instance(Device::DeviceType::kCUDA); + pg_factory->GetOrCreate(nn::parallel::GetDataParallelProcessGroupName(rank.GlobalRank()), + nn::parallel::GetDataParallelGroupRanks(rank.GlobalRank())); + auto model = std::make_shared(4, 4, /*bias=*/false, + Device(Device::DeviceType::kCUDA, + nn::parallel::global::GetLocalProcRank())); + nn::parallel::DistributedDataParallelConfig ddp_config; + ddp_config.zero_stage = 1; + ddp_config.overlap_grad_reduce = false; + ddp_config.overlap_param_gather = false; + auto ddp_model = std::make_shared(model, rank, ddp_config); + nn::parallel::DistributedOptimizer optimizer(optimizers::SGD::CreateNamed(0.1f), model->NamedParameters(), + std::vector>{ddp_model}, + /*ddp_world_size=*/2, /*ddp_rank=*/rank.GlobalRank()); + auto parameter = model->Parameters().front(); + ASSERT_NE(parameter->grad(), nullptr); + parameter->grad()->Fill(1.0f); + auto total_norm = optimizer.ClipGradNorm(model->Parameters(), 2.0f, 2.0f); + auto total_norm_cpu = total_norm->To(Device()); + EXPECT_NEAR(*static_cast(total_norm_cpu.DataPtr()), 4.0f, 1e-4f); +} + INFINI_TRAIN_REGISTER_TEST(OptimizerParameterNamesTest); + +TEST_P(OptimizerParameterNamesTest, DistributedOptimizerClipGradNormUsesGlobalShardNorm) { + ONLY_CUDA(); + REQUIRE_MIN_DEVICES(2); + if (nn::parallel::global::GetDataParallelSize() != 2) { + GTEST_SKIP() << "requires PROC_WORLD_SIZE=2"; + } + + const nn::parallel::Rank rank(nn::parallel::global::GetGlobalProcRank(), /*thread_rank=*/0, + nn::parallel::global::GetNprocPerNode(), /*thread_size=*/1); + auto *pg_factory = nn::parallel::ProcessGroupFactory::Instance(Device::DeviceType::kCUDA); + pg_factory->GetOrCreate(nn::parallel::GetDataParallelProcessGroupName(rank.GlobalRank()), + nn::parallel::GetDataParallelGroupRanks(rank.GlobalRank())); + + auto model = std::make_shared(64, 4, /*bias=*/false, + Device(Device::DeviceType::kCUDA, + nn::parallel::global::GetLocalProcRank())); + nn::parallel::DistributedDataParallelConfig ddp_config; + ddp_config.zero_stage = 1; + ddp_config.overlap_grad_reduce = false; + ddp_config.overlap_param_gather = false; + auto ddp_model = std::make_shared(model, rank, ddp_config); + const auto named_parameters = model->NamedParameters(); + + nn::parallel::DistributedOptimizer optimizer(optimizers::SGD::CreateNamed(0.1f), named_parameters, + std::vector>{ddp_model}, + /*ddp_world_size=*/2, /*ddp_rank=*/rank.GlobalRank()); + for (const auto ¶meter : model->Parameters()) { + ASSERT_NE(parameter->grad(), nullptr); + parameter->grad()->Fill(1.0f); + } + + auto total_norm = optimizer.ClipGradNorm(model->Parameters(), 2.0f, 2.0f); + auto total_norm_cpu = total_norm->To(Device()); + EXPECT_NEAR(*static_cast(total_norm_cpu.DataPtr()), 16.0f, 1e-4f); +} From 2dc16adf6a1455a754862b4cca351e380d8f8617 Mon Sep 17 00:00:00 2001 From: zhengjiahong Date: Mon, 14 Sep 2026 10:12:31 +0800 Subject: [PATCH 2/7] feat: integrate gradient clipping into training examples --- example/gpt2/main.cc | 15 +++++++++++++++ example/llama3/main.cc | 15 +++++++++++++++ .../src/nn/parallel/pp/pipeline_schedule.cc | 2 ++ 3 files changed, 32 insertions(+) diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc index 8e5d92c02..d7a257366 100644 --- a/example/gpt2/main.cc +++ b/example/gpt2/main.cc @@ -64,6 +64,10 @@ DEFINE_uint32(text_length, 64, "the length of the generated text"); // optimization DEFINE_double(learning_rate, 1e-4, "Peak learning rate."); DEFINE_int32(zero_stage, 0, "ZeRO stage (0/1/2/3); 0 disables DistributedOptimizer"); +DEFINE_double(clip_grad_norm, -1.0, "Maximum gradient norm; negative disables clipping."); +DEFINE_double(grad_norm_type, 2.0, "Gradient norm type (positive finite p or inf)."); +DEFINE_bool(clip_grad_error_if_nonfinite, true, "Fail if the pre-clipping gradient norm is NaN or Inf."); +DEFINE_string(clip_grad_foreach, "auto", "Gradient clipping path: auto|true|false."); // lr scheduler DEFINE_double(min_lr, 0.0, "Minimum learning rate."); DEFINE_string(lr_decay_style, "constant", "LR decay style: none|constant|linear|cosine|inverse-square-root"); @@ -352,6 +356,16 @@ void Train(const nn::parallel::Rank &rank) { optimizer = optimizer_creator(named_parameters); } + if (FLAGS_clip_grad_norm >= 0.0) { + std::optional foreach = std::nullopt; + if (FLAGS_clip_grad_foreach == "true") foreach = true; + else if (FLAGS_clip_grad_foreach == "false") foreach = false; + else CHECK_EQ(FLAGS_clip_grad_foreach, "auto"); + optimizer->SetClipGradNormConfig(static_cast(FLAGS_clip_grad_norm), + static_cast(FLAGS_grad_norm_type), + FLAGS_clip_grad_error_if_nonfinite, foreach); + } + const int64_t lr_decay_iters = FLAGS_lr_decay_iters > 0 ? FLAGS_lr_decay_iters : FLAGS_num_iteration; TrainingLRSchedulerConfig sched_config; sched_config.lr = static_cast(FLAGS_learning_rate); @@ -510,6 +524,7 @@ void Train(const nn::parallel::Rank &rank) { LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish backward"; } + optimizer->ClipGradNormConfigured(); optimizer->Step(); if (scheduler) { scheduler->Step(); diff --git a/example/llama3/main.cc b/example/llama3/main.cc index 19620e993..3fe7256d8 100644 --- a/example/llama3/main.cc +++ b/example/llama3/main.cc @@ -63,6 +63,10 @@ DEFINE_uint32(text_length, 64, "the length of the generated text"); // optimization DEFINE_double(learning_rate, 1e-5, "Peak learning rate."); DEFINE_int32(zero_stage, 0, "ZeRO stage (0/1/2/3); 0 disables DistributedOptimizer"); +DEFINE_double(clip_grad_norm, -1.0, "Maximum gradient norm; negative disables clipping."); +DEFINE_double(grad_norm_type, 2.0, "Gradient norm type (positive finite p or inf)."); +DEFINE_bool(clip_grad_error_if_nonfinite, true, "Fail if the pre-clipping gradient norm is NaN or Inf."); +DEFINE_string(clip_grad_foreach, "auto", "Gradient clipping path: auto|true|false."); // lr scheduler DEFINE_double(min_lr, 0.0, "Minimum learning rate."); DEFINE_string(lr_decay_style, "constant", "LR decay style: none|constant|linear|cosine|inverse-square-root"); @@ -334,6 +338,16 @@ void Train(const nn::parallel::Rank &rank) { optimizer = optimizer_creator(named_parameters); } + if (FLAGS_clip_grad_norm >= 0.0) { + std::optional foreach = std::nullopt; + if (FLAGS_clip_grad_foreach == "true") foreach = true; + else if (FLAGS_clip_grad_foreach == "false") foreach = false; + else CHECK_EQ(FLAGS_clip_grad_foreach, "auto"); + optimizer->SetClipGradNormConfig(static_cast(FLAGS_clip_grad_norm), + static_cast(FLAGS_grad_norm_type), + FLAGS_clip_grad_error_if_nonfinite, foreach); + } + const int64_t lr_decay_iters = FLAGS_lr_decay_iters > 0 ? FLAGS_lr_decay_iters : FLAGS_num_iteration; TrainingLRSchedulerConfig sched_config; sched_config.lr = static_cast(FLAGS_learning_rate); @@ -489,6 +503,7 @@ void Train(const nn::parallel::Rank &rank) { LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish backward"; } + optimizer->ClipGradNormConfigured(); optimizer->Step(); if (scheduler) { scheduler->Step(); diff --git a/infini_train/src/nn/parallel/pp/pipeline_schedule.cc b/infini_train/src/nn/parallel/pp/pipeline_schedule.cc index 6578e628b..7a507d1b8 100644 --- a/infini_train/src/nn/parallel/pp/pipeline_schedule.cc +++ b/infini_train/src/nn/parallel/pp/pipeline_schedule.cc @@ -299,6 +299,8 @@ float PipelineSchedule::Step(std::shared_ptr input, std::shared_ptrClipGradNormConfigured(); optimizer->Step(); return lossf; From e2a845aece57b563e4899e959142eee6969d480a Mon Sep 17 00:00:00 2001 From: zhengjiahong Date: Mon, 14 Sep 2026 10:50:45 +0800 Subject: [PATCH 3/7] style: format gradient clipping changes --- example/gpt2/main.cc | 14 +++++---- example/llama3/main.cc | 14 +++++---- .../nn/parallel/ddp/distributed_optimizer.h | 6 ++-- infini_train/include/optimizer.h | 7 +++-- .../src/kernels/cpu/accumulate_grad.cc | 7 ++--- .../src/kernels/cuda/accumulate_grad.cu | 7 ++--- .../nn/parallel/ddp/distributed_optimizer.cc | 30 +++++++++---------- infini_train/src/optimizer.cc | 28 +++++++++-------- tests/optimizer/test_clip_grad_norm.cc | 17 ++++------- .../test_optimizer_parameter_names.cc | 22 +++++--------- 10 files changed, 75 insertions(+), 77 deletions(-) diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc index d7a257366..a4eeba12e 100644 --- a/example/gpt2/main.cc +++ b/example/gpt2/main.cc @@ -358,12 +358,16 @@ void Train(const nn::parallel::Rank &rank) { if (FLAGS_clip_grad_norm >= 0.0) { std::optional foreach = std::nullopt; - if (FLAGS_clip_grad_foreach == "true") foreach = true; - else if (FLAGS_clip_grad_foreach == "false") foreach = false; - else CHECK_EQ(FLAGS_clip_grad_foreach, "auto"); + if (FLAGS_clip_grad_foreach == "true") { + foreach { = true; } + } else if (FLAGS_clip_grad_foreach == "false") { + foreach { = false; } + } else { + CHECK_EQ(FLAGS_clip_grad_foreach, "auto"); + } optimizer->SetClipGradNormConfig(static_cast(FLAGS_clip_grad_norm), - static_cast(FLAGS_grad_norm_type), - FLAGS_clip_grad_error_if_nonfinite, foreach); + static_cast(FLAGS_grad_norm_type), FLAGS_clip_grad_error_if_nonfinite, + foreach); } const int64_t lr_decay_iters = FLAGS_lr_decay_iters > 0 ? FLAGS_lr_decay_iters : FLAGS_num_iteration; diff --git a/example/llama3/main.cc b/example/llama3/main.cc index 3fe7256d8..245133cbd 100644 --- a/example/llama3/main.cc +++ b/example/llama3/main.cc @@ -340,12 +340,16 @@ void Train(const nn::parallel::Rank &rank) { if (FLAGS_clip_grad_norm >= 0.0) { std::optional foreach = std::nullopt; - if (FLAGS_clip_grad_foreach == "true") foreach = true; - else if (FLAGS_clip_grad_foreach == "false") foreach = false; - else CHECK_EQ(FLAGS_clip_grad_foreach, "auto"); + if (FLAGS_clip_grad_foreach == "true") { + foreach { = true; } + } else if (FLAGS_clip_grad_foreach == "false") { + foreach { = false; } + } else { + CHECK_EQ(FLAGS_clip_grad_foreach, "auto"); + } optimizer->SetClipGradNormConfig(static_cast(FLAGS_clip_grad_norm), - static_cast(FLAGS_grad_norm_type), - FLAGS_clip_grad_error_if_nonfinite, foreach); + static_cast(FLAGS_grad_norm_type), FLAGS_clip_grad_error_if_nonfinite, + foreach); } const int64_t lr_decay_iters = FLAGS_lr_decay_iters > 0 ? FLAGS_lr_decay_iters : FLAGS_num_iteration; diff --git a/infini_train/include/nn/parallel/ddp/distributed_optimizer.h b/infini_train/include/nn/parallel/ddp/distributed_optimizer.h index 7edacb3eb..52e86db5d 100644 --- a/infini_train/include/nn/parallel/ddp/distributed_optimizer.h +++ b/infini_train/include/nn/parallel/ddp/distributed_optimizer.h @@ -33,9 +33,9 @@ class DistributedOptimizer final : public infini_train::Optimizer { void Step() override; - std::shared_ptr ClipGradNorm_( - const std::vector> ¶meters, float max_norm, float norm_type = 2.0f, - bool error_if_nonfinite = false, std::optional foreach = std::nullopt) override; + std::shared_ptr ClipGradNorm_(const std::vector> ¶meters, float max_norm, + float norm_type = 2.0f, bool error_if_nonfinite = false, + std::optional foreach = std::nullopt) override; void ZeroGrad(bool set_to_none = true) override; diff --git a/infini_train/include/optimizer.h b/infini_train/include/optimizer.h index 90d49f430..5b3007815 100644 --- a/infini_train/include/optimizer.h +++ b/infini_train/include/optimizer.h @@ -29,9 +29,10 @@ class Optimizer { virtual void ZeroGrad(bool set_to_none = true); // Return the pre-clipping norm and scale selected gradients in place. - virtual std::shared_ptr ClipGradNorm_( - const std::vector> ¶meters, float max_norm, float norm_type = 2.0f, - bool error_if_nonfinite = false, std::optional foreach = std::nullopt); + virtual std::shared_ptr ClipGradNorm_(const std::vector> ¶meters, + float max_norm, float norm_type = 2.0f, + bool error_if_nonfinite = false, + std::optional foreach = std::nullopt); std::shared_ptr ClipGradNorm(const std::vector> ¶meters, float max_norm, float norm_type = 2.0f, bool error_if_nonfinite = false, diff --git a/infini_train/src/kernels/cpu/accumulate_grad.cc b/infini_train/src/kernels/cpu/accumulate_grad.cc index a78b5175b..25cf525e7 100644 --- a/infini_train/src/kernels/cpu/accumulate_grad.cc +++ b/infini_train/src/kernels/cpu/accumulate_grad.cc @@ -6,13 +6,10 @@ namespace infini_train::kernels::cpu { -template -void ScaleInplaceTyped(const std::shared_ptr &tensor, float scale) { +template void ScaleInplaceTyped(const std::shared_ptr &tensor, float scale) { auto *data = static_cast(tensor->DataPtr()); #pragma omp parallel for - for (size_t i = 0; i < tensor->NumElements(); ++i) { - data[i] = T(static_cast(data[i]) * scale); - } + for (size_t i = 0; i < tensor->NumElements(); ++i) { data[i] = T(static_cast(data[i]) * scale); } } void ScaleInplace(const std::shared_ptr &tensor, float scale) { diff --git a/infini_train/src/kernels/cuda/accumulate_grad.cu b/infini_train/src/kernels/cuda/accumulate_grad.cu index b107b37ef..808a51b89 100644 --- a/infini_train/src/kernels/cuda/accumulate_grad.cu +++ b/infini_train/src/kernels/cuda/accumulate_grad.cu @@ -11,8 +11,7 @@ namespace infini_train::kernels::cuda { -template -__global__ void ScaleInplaceKernel(T *data, float scale, size_t num_elements) { +template __global__ void ScaleInplaceKernel(T *data, float scale, size_t num_elements) { const size_t idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < num_elements) { data[idx] = common::cuda::Mul(data[idx], common::cuda::Cast(scale)); @@ -30,8 +29,8 @@ void ScaleInplace(const std::shared_ptr &tensor, float scale) { core::cuda::DispatchCudaFunc( tensor->Dtype(), [=]() { - ScaleInplaceKernel<<>>( - static_cast(tensor->DataPtr()), scale, num_elements); + ScaleInplaceKernel<<>>(static_cast(tensor->DataPtr()), + scale, num_elements); }, "CUDA ScaleInplace"); } diff --git a/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc b/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc index afb10649c..11dda6e7b 100644 --- a/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc +++ b/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc @@ -46,7 +46,7 @@ DistributedOptimizer::DistributedOptimizer(OptimizerCreatorNamed creator, const NamedParameterList shard_named_parameters; BuildShardParamsAndBindGrads( [this, ¶meter_name_by_tensor, &shard_named_parameters](const std::shared_ptr ¶meter, - const std::shared_ptr ¶m_piece) { + const std::shared_ptr ¶m_piece) { const auto name_it = parameter_name_by_tensor.find(parameter.get()); CHECK(name_it != parameter_name_by_tensor.end()) << "DistributedOptimizer parameter is not registered in the model"; @@ -161,9 +161,9 @@ void DistributedOptimizer::FinishParamSync(bool skip_next_bucket_dispatch) { for (auto &group : bucket_groups_) { group->FinishParamSync(skip_next_bucket_dispatch); } } -std::shared_ptr DistributedOptimizer::ClipGradNorm_( - const std::vector> ¶meters, float max_norm, float norm_type, - bool error_if_nonfinite, std::optional foreach) { +std::shared_ptr DistributedOptimizer::ClipGradNorm_(const std::vector> ¶meters, + float max_norm, float norm_type, bool error_if_nonfinite, + std::optional foreach) { CHECK_GE(max_norm, 0.0f) << "max_norm must be non-negative."; CHECK((norm_type > 0.0f && std::isfinite(norm_type)) || norm_type == std::numeric_limits::infinity()) << "norm_type must be positive finite or +inf."; @@ -185,11 +185,11 @@ std::shared_ptr DistributedOptimizer::ClipGradNorm_( if (requested.empty() || requested.contains(shard_param_owners_[i].get())) { selected_shards.push_back(shard_params_[i]); bool include_in_norm = true; - if (global::GetTensorParallelSize() > 1 && - tp_replicated_param_owners_.contains(shard_param_owners_[i].get())) { - const auto *tp_group = ProcessGroupFactory::Instance(shard_params_[i]->GetDevice().type()) - ->Get(GetTensorParallelProcessGroupName( - shard_params_[i]->GetDevice().Rank().GlobalRank())); + if (global::GetTensorParallelSize() > 1 + && tp_replicated_param_owners_.contains(shard_param_owners_[i].get())) { + const auto *tp_group + = ProcessGroupFactory::Instance(shard_params_[i]->GetDevice().type()) + ->Get(GetTensorParallelProcessGroupName(shard_params_[i]->GetDevice().Rank().GlobalRank())); CHECK(tp_group) << "Tensor-parallel process group is not initialized."; include_in_norm = tp_group->GetGroupRank(shard_params_[i]->GetDevice().Rank().GlobalRank()) == 0; } @@ -201,8 +201,8 @@ std::shared_ptr DistributedOptimizer::ClipGradNorm_( // Ask the base optimizer for the local norm without scaling. Infinite // max_norm makes its coefficient exactly one for finite gradients. - auto local_norm_tensor = base_optimizer_->ClipGradNorm_( - norm_shards, std::numeric_limits::infinity(), norm_type, false, std::nullopt); + auto local_norm_tensor = base_optimizer_->ClipGradNorm_(norm_shards, std::numeric_limits::infinity(), + norm_type, false, std::nullopt); const float local_norm = *static_cast(local_norm_tensor->DataPtr()); const bool is_inf_norm = std::isinf(norm_type); double local_stat = is_inf_norm ? static_cast(local_norm) @@ -230,14 +230,14 @@ std::shared_ptr DistributedOptimizer::ClipGradNorm_( group->AllReduce(reduced, is_inf_norm ? function::ReduceOpType::kMax : function::ReduceOpType::kSum, false); } if (global::GetTensorParallelSize() > 1) { - const auto *tp_group = ProcessGroupFactory::Instance(total_norm_device.type())->Get( - GetTensorParallelProcessGroupName(total_norm_device.Rank().GlobalRank())); + const auto *tp_group = ProcessGroupFactory::Instance(total_norm_device.type()) + ->Get(GetTensorParallelProcessGroupName(total_norm_device.Rank().GlobalRank())); CHECK(tp_group) << "Tensor-parallel process group is not initialized."; tp_group->AllReduce(reduced, is_inf_norm ? function::ReduceOpType::kMax : function::ReduceOpType::kSum, false); } if (nn::parallel::global::GetPipelineParallelSize() > 1) { - const auto *pp_group = ProcessGroupFactory::Instance(total_norm_device.type())->Get( - GetPipelineParallelProcessGroupName(total_norm_device.Rank().GlobalRank())); + const auto *pp_group = ProcessGroupFactory::Instance(total_norm_device.type()) + ->Get(GetPipelineParallelProcessGroupName(total_norm_device.Rank().GlobalRank())); CHECK(pp_group) << "Pipeline process group is not initialized."; pp_group->AllReduce(reduced, is_inf_norm ? function::ReduceOpType::kMax : function::ReduceOpType::kSum, false); } diff --git a/infini_train/src/optimizer.cc b/infini_train/src/optimizer.cc index f57219537..26d922263 100644 --- a/infini_train/src/optimizer.cc +++ b/infini_train/src/optimizer.cc @@ -96,7 +96,7 @@ void ScaleGradientInplace(const std::shared_ptr &gradient, float scale) } // namespace void Optimizer::SetClipGradNormConfig(float max_norm, float norm_type, bool error_if_nonfinite, - std::optional foreach) { + std::optional foreach) { clip_grad_norm_config_ = ClipGradNormConfig{max_norm, norm_type, error_if_nonfinite, foreach}; } @@ -118,8 +118,8 @@ void Optimizer::ScaleGradients_(const std::vector> ¶ } } -std::shared_ptr Optimizer::ClipGradNorm_(const std::vector> ¶meters, - float max_norm, float norm_type, bool error_if_nonfinite, +std::shared_ptr Optimizer::ClipGradNorm_(const std::vector> ¶meters, float max_norm, + float norm_type, bool error_if_nonfinite, std::optional foreach) { CHECK_GE(max_norm, 0.0f) << "max_norm must be non-negative."; CHECK((norm_type > 0.0f && std::isfinite(norm_type)) || norm_type == std::numeric_limits::infinity()) @@ -139,18 +139,23 @@ std::shared_ptr Optimizer::ClipGradNorm_(const std::vector 1) { const auto device = [&]() { for (const auto ¶meter : parameters) { - if (parameter && parameter->grad()) return parameter->grad()->GetDevice(); + if (parameter && parameter->grad()) { + return parameter->grad()->GetDevice(); + } } return Device(); }(); auto reduced = std::make_shared(std::vector{}, DataType::kFLOAT32, device); reduced->Fill(static_cast(std::isinf(norm_type) ? stats.max_abs : stats.sum)); - const auto *group = nn::parallel::ProcessGroupFactory::Instance(device.type())->Get( - nn::parallel::GetPipelineParallelProcessGroupName(device.Rank().GlobalRank())); + const auto *group = nn::parallel::ProcessGroupFactory::Instance(device.type()) + ->Get(nn::parallel::GetPipelineParallelProcessGroupName(device.Rank().GlobalRank())); CHECK(group) << "Pipeline process group is not initialized."; - group->AllReduce(reduced, std::isinf(norm_type) ? nn::parallel::function::ReduceOpType::kMax - : nn::parallel::function::ReduceOpType::kSum, false); - Tensor reduced_cpu = reduced->GetDevice().IsCPU() ? Tensor(*reduced, 0, reduced->Dims()) : reduced->To(Device()); + group->AllReduce(reduced, + std::isinf(norm_type) ? nn::parallel::function::ReduceOpType::kMax + : nn::parallel::function::ReduceOpType::kSum, + false); + Tensor reduced_cpu + = reduced->GetDevice().IsCPU() ? Tensor(*reduced, 0, reduced->Dims()) : reduced->To(Device()); const float value = *static_cast(reduced_cpu.DataPtr()); total_norm = std::isinf(norm_type) ? value : std::pow(static_cast(value), 1.0 / norm_type); } @@ -158,9 +163,8 @@ std::shared_ptr Optimizer::ClipGradNorm_(const std::vector(max_norm) / (total_norm + 1e-6), 1.0); + const double coefficient + = std::isinf(max_norm) ? 1.0 : std::min(static_cast(max_norm) / (total_norm + 1e-6), 1.0); if (stats.has_gradient) { ScaleGradients_(parameters, static_cast(coefficient)); } diff --git a/tests/optimizer/test_clip_grad_norm.cc b/tests/optimizer/test_clip_grad_norm.cc index 6a690258d..92e4ffa7c 100644 --- a/tests/optimizer/test_clip_grad_norm.cc +++ b/tests/optimizer/test_clip_grad_norm.cc @@ -21,7 +21,7 @@ static std::shared_ptr MakeTensor(Device device, const std::vector(tensor->DataPtr())); } else { auto cpu = std::make_shared(std::vector{static_cast(values.size())}, - DataType::kFLOAT32, Device()); + DataType::kFLOAT32, Device()); std::copy(values.begin(), values.end(), static_cast(cpu->DataPtr())); tensor->CopyFrom(*cpu); } @@ -37,8 +37,7 @@ TEST_P(ClipGradNormTest, L2ClipsInPlace) { auto param = std::make_shared(std::vector{2}, DataType::kFLOAT32, GetDevice()); auto grad = MakeTensor(GetDevice(), {3.0f, 4.0f}); param->set_grad(grad); - auto optimizer = std::make_shared( - std::vector>{param}, 0.1f); + auto optimizer = std::make_shared(std::vector>{param}, 0.1f); auto total_norm = optimizer->ClipGradNorm({param}, 2.0f, 2.0f); EXPECT_NEAR(ScalarCPU(total_norm), 5.0f, 1e-5f); @@ -52,8 +51,7 @@ TEST_P(ClipGradNormTest, SupportsL1AndInfinity) { auto param = std::make_shared(std::vector{3}, DataType::kFLOAT32, GetDevice()); auto grad = MakeTensor(GetDevice(), {-2.0f, 1.0f, 3.0f}); param->set_grad(grad); - auto optimizer = std::make_shared( - std::vector>{param}, 0.1f); + auto optimizer = std::make_shared(std::vector>{param}, 0.1f); auto l1 = optimizer->ClipGradNorm({param}, 3.0f, 1.0f); EXPECT_NEAR(ScalarCPU(l1), 6.0f, 1e-5f); @@ -77,8 +75,7 @@ TEST_P(ClipGradNormTest, IgnoresMissingAndDuplicateGradients) { auto first = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice()); first->set_grad(MakeTensor(GetDevice(), {3.0f})); auto second = std::make_shared(std::vector{1}, DataType::kFLOAT32, GetDevice()); - auto optimizer = std::make_shared( - std::vector>{first, second}, 0.1f); + auto optimizer = std::make_shared(std::vector>{first, second}, 0.1f); auto total_norm = optimizer->ClipGradNorm({first, first, second}, 1.0f, 2.0f); EXPECT_NEAR(ScalarCPU(total_norm), 3.0f, 1e-5f); @@ -95,8 +92,7 @@ TEST_P(ClipGradNormTest, NonIntegerPAndZeroMaxNorm) { auto param = std::make_shared(std::vector{2}, DataType::kFLOAT32, GetDevice()); auto grad = MakeTensor(GetDevice(), {2.0f, 2.0f}); param->set_grad(grad); - auto optimizer = std::make_shared( - std::vector>{param}, 0.1f); + auto optimizer = std::make_shared(std::vector>{param}, 0.1f); const float expected = std::pow(2.0f * std::pow(2.0f, 3.5f), 1.0f / 3.5f); auto total_norm = optimizer->ClipGradNorm({param}, 0.0f, 3.5f); @@ -112,7 +108,6 @@ TEST_P(ClipGradNormTest, ErrorOnNonFiniteBeforeScaling) { auto param = std::make_shared(std::vector{2}, DataType::kFLOAT32, GetDevice()); auto grad = MakeTensor(GetDevice(), {std::numeric_limits::quiet_NaN(), 1.0f}); param->set_grad(grad); - auto optimizer = std::make_shared( - std::vector>{param}, 0.1f); + auto optimizer = std::make_shared(std::vector>{param}, 0.1f); EXPECT_DEATH(optimizer->ClipGradNorm({param}, 1.0f, 2.0f, true), "non-finite"); } diff --git a/tests/optimizer/test_optimizer_parameter_names.cc b/tests/optimizer/test_optimizer_parameter_names.cc index f26ede2f2..c96f5ca4b 100644 --- a/tests/optimizer/test_optimizer_parameter_names.cc +++ b/tests/optimizer/test_optimizer_parameter_names.cc @@ -65,9 +65,8 @@ TEST_P(OptimizerParameterNamesTest, DistributedOptimizerPropagatesNamesToShardOp pg_factory->GetOrCreate(nn::parallel::GetDataParallelProcessGroupName(rank.GlobalRank()), nn::parallel::GetDataParallelGroupRanks(rank.GlobalRank())); - auto model = std::make_shared(64, 4, /*bias=*/false, - Device(Device::DeviceType::kCUDA, - nn::parallel::global::GetLocalProcRank())); + auto model = std::make_shared( + 64, 4, /*bias=*/false, Device(Device::DeviceType::kCUDA, nn::parallel::global::GetLocalProcRank())); const auto named_parameters = model->NamedParameters(); nn::parallel::DistributedDataParallelConfig ddp_config; @@ -97,7 +96,6 @@ TEST_P(OptimizerParameterNamesTest, PreservesNumericKeysWhenNamesAreNotSet) { EXPECT_TRUE(state.contains("adam.v.0")); } - TEST_P(OptimizerParameterNamesTest, DistributedOptimizerClipGradNormUsesZero2LocalShard) { ONLY_CUDA(); REQUIRE_MIN_DEVICES(2); @@ -109,9 +107,8 @@ TEST_P(OptimizerParameterNamesTest, DistributedOptimizerClipGradNormUsesZero2Loc auto *pg_factory = nn::parallel::ProcessGroupFactory::Instance(Device::DeviceType::kCUDA); pg_factory->GetOrCreate(nn::parallel::GetDataParallelProcessGroupName(rank.GlobalRank()), nn::parallel::GetDataParallelGroupRanks(rank.GlobalRank())); - auto model = std::make_shared(64, 4, /*bias=*/false, - Device(Device::DeviceType::kCUDA, - nn::parallel::global::GetLocalProcRank())); + auto model = std::make_shared( + 64, 4, /*bias=*/false, Device(Device::DeviceType::kCUDA, nn::parallel::global::GetLocalProcRank())); nn::parallel::DistributedDataParallelConfig ddp_config; ddp_config.zero_stage = 2; ddp_config.overlap_grad_reduce = false; @@ -133,7 +130,6 @@ TEST_P(OptimizerParameterNamesTest, DistributedOptimizerClipGradNormUsesZero2Loc EXPECT_NEAR(static_cast(local_shard.DataPtr())[0], 0.5f, 1e-4f); } - TEST_P(OptimizerParameterNamesTest, DistributedOptimizerClipGradNormHandlesEmptyLocalShard) { ONLY_CUDA(); REQUIRE_MIN_DEVICES(2); @@ -145,9 +141,8 @@ TEST_P(OptimizerParameterNamesTest, DistributedOptimizerClipGradNormHandlesEmpty auto *pg_factory = nn::parallel::ProcessGroupFactory::Instance(Device::DeviceType::kCUDA); pg_factory->GetOrCreate(nn::parallel::GetDataParallelProcessGroupName(rank.GlobalRank()), nn::parallel::GetDataParallelGroupRanks(rank.GlobalRank())); - auto model = std::make_shared(4, 4, /*bias=*/false, - Device(Device::DeviceType::kCUDA, - nn::parallel::global::GetLocalProcRank())); + auto model = std::make_shared( + 4, 4, /*bias=*/false, Device(Device::DeviceType::kCUDA, nn::parallel::global::GetLocalProcRank())); nn::parallel::DistributedDataParallelConfig ddp_config; ddp_config.zero_stage = 1; ddp_config.overlap_grad_reduce = false; @@ -179,9 +174,8 @@ TEST_P(OptimizerParameterNamesTest, DistributedOptimizerClipGradNormUsesGlobalSh pg_factory->GetOrCreate(nn::parallel::GetDataParallelProcessGroupName(rank.GlobalRank()), nn::parallel::GetDataParallelGroupRanks(rank.GlobalRank())); - auto model = std::make_shared(64, 4, /*bias=*/false, - Device(Device::DeviceType::kCUDA, - nn::parallel::global::GetLocalProcRank())); + auto model = std::make_shared( + 64, 4, /*bias=*/false, Device(Device::DeviceType::kCUDA, nn::parallel::global::GetLocalProcRank())); nn::parallel::DistributedDataParallelConfig ddp_config; ddp_config.zero_stage = 1; ddp_config.overlap_grad_reduce = false; From 7ada2df7674c2b009ab20ee88f61f480ccb73fae Mon Sep 17 00:00:00 2001 From: zhengjiahong Date: Mon, 14 Sep 2026 11:04:37 +0800 Subject: [PATCH 4/7] fix: preserve foreach option during formatting --- example/gpt2/main.cc | 8 ++++---- example/llama3/main.cc | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc index a4eeba12e..203d56115 100644 --- a/example/gpt2/main.cc +++ b/example/gpt2/main.cc @@ -357,17 +357,17 @@ void Train(const nn::parallel::Rank &rank) { } if (FLAGS_clip_grad_norm >= 0.0) { - std::optional foreach = std::nullopt; + std::optional foreach_option = std::nullopt; if (FLAGS_clip_grad_foreach == "true") { - foreach { = true; } + foreach_option = true; } else if (FLAGS_clip_grad_foreach == "false") { - foreach { = false; } + foreach_option = false; } else { CHECK_EQ(FLAGS_clip_grad_foreach, "auto"); } optimizer->SetClipGradNormConfig(static_cast(FLAGS_clip_grad_norm), static_cast(FLAGS_grad_norm_type), FLAGS_clip_grad_error_if_nonfinite, - foreach); + foreach_option); } const int64_t lr_decay_iters = FLAGS_lr_decay_iters > 0 ? FLAGS_lr_decay_iters : FLAGS_num_iteration; diff --git a/example/llama3/main.cc b/example/llama3/main.cc index 245133cbd..935623f18 100644 --- a/example/llama3/main.cc +++ b/example/llama3/main.cc @@ -339,17 +339,17 @@ void Train(const nn::parallel::Rank &rank) { } if (FLAGS_clip_grad_norm >= 0.0) { - std::optional foreach = std::nullopt; + std::optional foreach_option = std::nullopt; if (FLAGS_clip_grad_foreach == "true") { - foreach { = true; } + foreach_option = true; } else if (FLAGS_clip_grad_foreach == "false") { - foreach { = false; } + foreach_option = false; } else { CHECK_EQ(FLAGS_clip_grad_foreach, "auto"); } optimizer->SetClipGradNormConfig(static_cast(FLAGS_clip_grad_norm), static_cast(FLAGS_grad_norm_type), FLAGS_clip_grad_error_if_nonfinite, - foreach); + foreach_option); } const int64_t lr_decay_iters = FLAGS_lr_decay_iters > 0 ? FLAGS_lr_decay_iters : FLAGS_num_iteration; From d53d0a4ba30ad18fdf36628950b7ad5a3f7ce587 Mon Sep 17 00:00:00 2001 From: zhengjiahong Date: Tue, 15 Sep 2026 09:06:35 +0800 Subject: [PATCH 5/7] fix: report clipped gradient norm in examples --- example/gpt2/main.cc | 23 +++++++++++++++++------ example/llama3/main.cc | 23 +++++++++++++++++------ 2 files changed, 34 insertions(+), 12 deletions(-) diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc index 203d56115..da4d8da0f 100644 --- a/example/gpt2/main.cc +++ b/example/gpt2/main.cc @@ -483,6 +483,7 @@ void Train(const nn::parallel::Rank &rank) { const float current_lr = scheduler ? scheduler->learning_rate() : static_cast(FLAGS_learning_rate); float lossf = 0.0f; + std::optional total_grad_norm; // model->Train(); if (pp_world_size == 1) { optimizer->ZeroGrad(); @@ -528,7 +529,13 @@ void Train(const nn::parallel::Rank &rank) { LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish backward"; } - optimizer->ClipGradNormConfigured(); + if (optimizer->HasClipGradNormConfig()) { + auto norm_tensor = optimizer->ClipGradNormConfigured(); + if (norm_tensor) { + auto total_grad_norm_cpu = norm_tensor->To(Device()); + total_grad_norm = *static_cast(total_grad_norm_cpu.DataPtr()); + } + } optimizer->Step(); if (scheduler) { scheduler->Step(); @@ -557,11 +564,15 @@ void Train(const nn::parallel::Rank &rank) { if (rank.IsLastRank()) { size_t used_mb = 0, reserved_mb = 0; std::tie(used_mb, reserved_mb) = impl->GetMemPoolPeakMB(device); - LOG(ERROR) << std::format("step {:4d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} tok/s | " - "peak used: {:5d} MB | peak reserved: {:5d} MB, DP={}, TP={}, SP={}, PP={})", - step + 1, FLAGS_num_iteration, lossf, current_lr, duration_us / 1e3f, tps, - used_mb, reserved_mb, ddp_world_size, tp_world_size, sp_world_size, - pp_world_size); + auto message = std::format( + "step {:4d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} tok/s | " + "peak used: {:5d} MB | peak reserved: {:5d} MB, DP={}, TP={}, SP={}, PP={})", + step + 1, FLAGS_num_iteration, lossf, current_lr, duration_us / 1e3f, tps, used_mb, reserved_mb, + ddp_world_size, tp_world_size, sp_world_size, pp_world_size); + if (total_grad_norm) { + message += std::format(" | total_grad_norm {:.6f}", *total_grad_norm); + } + LOG(ERROR) << message; if ((step + 1) % FLAGS_freq_generate_txt == 0) { if (tokenizer) { diff --git a/example/llama3/main.cc b/example/llama3/main.cc index 935623f18..c70cd9274 100644 --- a/example/llama3/main.cc +++ b/example/llama3/main.cc @@ -463,6 +463,7 @@ void Train(const nn::parallel::Rank &rank) { const float current_lr = scheduler ? scheduler->learning_rate() : static_cast(FLAGS_learning_rate); float lossf = 0.0f; + std::optional total_grad_norm; if (pp_world_size == 1) { // model->Train(); optimizer->ZeroGrad(); @@ -507,7 +508,13 @@ void Train(const nn::parallel::Rank &rank) { LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish backward"; } - optimizer->ClipGradNormConfigured(); + if (optimizer->HasClipGradNormConfig()) { + auto norm_tensor = optimizer->ClipGradNormConfigured(); + if (norm_tensor) { + auto total_grad_norm_cpu = norm_tensor->To(Device()); + total_grad_norm = *static_cast(total_grad_norm_cpu.DataPtr()); + } + } optimizer->Step(); if (scheduler) { scheduler->Step(); @@ -536,11 +543,15 @@ void Train(const nn::parallel::Rank &rank) { if (rank.IsLastRank()) { size_t used_mb = 0, reserved_mb = 0; std::tie(used_mb, reserved_mb) = impl->GetMemPoolPeakMB(device); - LOG(ERROR) << std::format("step {:4d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} tok/s | " - "peak used: {:5d} MB | peak reserved: {:5d} MB, DP={}, TP={}, SP={}, PP={})", - step + 1, FLAGS_num_iteration, lossf, current_lr, duration_us / 1e3f, tps, - used_mb, reserved_mb, ddp_world_size, tp_world_size, sp_world_size, - pp_world_size); + auto message = std::format( + "step {:4d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} tok/s | " + "peak used: {:5d} MB | peak reserved: {:5d} MB, DP={}, TP={}, SP={}, PP={})", + step + 1, FLAGS_num_iteration, lossf, current_lr, duration_us / 1e3f, tps, used_mb, reserved_mb, + ddp_world_size, tp_world_size, sp_world_size, pp_world_size); + if (total_grad_norm) { + message += std::format(" | total_grad_norm {:.6f}", *total_grad_norm); + } + LOG(ERROR) << message; if ((step + 1) % FLAGS_freq_generate_txt == 0) { // FIXME(jym): to support PP From b350e856f9f8bca36ccea84b7d70572e1dfd3c1d Mon Sep 17 00:00:00 2001 From: zhengjiahong Date: Tue, 15 Sep 2026 20:10:40 +0800 Subject: [PATCH 6/7] update --- docs/gradient_norm_clipping.md | 6 +- example/gpt2/main.cc | 10 +- example/llama3/main.cc | 10 +- .../nn/parallel/pp/pipeline_parallel.h | 2 + .../nn/parallel/pp/pipeline_schedule.h | 3 + infini_train/include/optimizer.h | 3 +- .../src/kernels/cpu/accumulate_grad.cc | 49 ++++++++ .../src/kernels/cuda/accumulate_grad.cu | 70 ++++++++++++ .../nn/parallel/ddp/distributed_optimizer.cc | 38 ++++--- .../src/nn/parallel/pp/pipeline_parallel.cc | 4 + .../src/nn/parallel/pp/pipeline_schedule.cc | 2 +- infini_train/src/optimizer.cc | 107 ++++++++++++++---- scripts/test_config.json | 70 +++++++++++- tests/optimizer/test_clip_grad_norm.cc | 73 ++++++++++++ .../test_optimizer_parameter_names.cc | 46 +++++++- 15 files changed, 447 insertions(+), 46 deletions(-) diff --git a/docs/gradient_norm_clipping.md b/docs/gradient_norm_clipping.md index 9966517f5..41ed90d0b 100644 --- a/docs/gradient_norm_clipping.md +++ b/docs/gradient_norm_clipping.md @@ -4,9 +4,9 @@ The public entry point is: auto total_norm = optimizer->ClipGradNorm_(parameters, max_norm, norm_type, error_if_nonfinite, std::nullopt); -It returns a CPU FP32 scalar containing the pre-clipping norm and scales the selected gradients in place. Parameters without gradients are ignored and duplicate parameter pointers are counted once. Supported norm types are positive finite p and positive infinity. The coefficient is min(max_norm / (total_norm + 1e-6), 1); max_norm == 0 therefore zeros gradients rather than disabling clipping. +It returns a CPU FP32 scalar containing the pre-clipping norm and scales the selected gradients in place. Parameters without gradients are ignored and duplicate parameter pointers are counted once. `norm_type` follows PyTorch-style vector norms: `0` first counts each gradient tensor that contains a non-zero element and then takes the 0-norm of that tensor list, finite positive and negative p use the p-norm formula, `+inf` takes the maximum absolute value, and `-inf` takes the minimum absolute value. The coefficient is min(max_norm / (total_norm + 1e-6), 1); max_norm == 0 therefore zeros gradients rather than disabling clipping. -foreach=true is intentionally rejected because a fused multi-tensor kernel is not part of this implementation yet. The CPU and CUDA ScaleInplace kernels write through the existing gradient view, preserving flat-buffer and ZeRO-shard aliases. CUDA norm inspection synchronizes only to obtain scalar statistics; scaling remains on the CUDA stream. +`foreach=true` selects the multi-tensor scaling path. CPU and CUDA `ScaleInplaceMulti` kernels write through the existing gradient views, preserving flat-buffer and ZeRO-shard aliases while issuing one batch dispatch per device/dtype group. CUDA norm inspection synchronizes only to obtain scalar statistics; scaling remains on the CUDA stream. DistributedOptimizer::ClipGradNorm_ first finishes pending gradient collectives, computes the norm over the shard parameters consumed by its base optimizer, reduces statistics over the DP process group, and applies the same coefficient to every local shard. When PP is enabled, statistics are also reduced over the PP process group. Pipeline scheduling calls ClipGradNormConfigured after all micro-batches complete. @@ -19,4 +19,4 @@ The GPT-2 and LLaMA-3 examples expose: A total_grad_norm log line is emitted when clipping is enabled. -Current limitation: tensor-parallel replicated-parameter ownership metadata is not yet represented by Tensor/Optimizer, so TP/SP duplicate filtering needs follow-up integration before claiming full mixed-parallel equivalence. +Tensor-parallel replicated-parameter ownership is filtered by DistributedOptimizer; complete mixed TP/SP/PP/vPP training coverage is provided by the optional matrix in `scripts/test_config.json` and should be run on a host with the corresponding checkpoints. diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc index da4d8da0f..36d5f1036 100644 --- a/example/gpt2/main.cc +++ b/example/gpt2/main.cc @@ -65,7 +65,7 @@ DEFINE_uint32(text_length, 64, "the length of the generated text"); DEFINE_double(learning_rate, 1e-4, "Peak learning rate."); DEFINE_int32(zero_stage, 0, "ZeRO stage (0/1/2/3); 0 disables DistributedOptimizer"); DEFINE_double(clip_grad_norm, -1.0, "Maximum gradient norm; negative disables clipping."); -DEFINE_double(grad_norm_type, 2.0, "Gradient norm type (positive finite p or inf)."); +DEFINE_double(grad_norm_type, 2.0, "Gradient norm type (finite p, inf, or -inf)."); DEFINE_bool(clip_grad_error_if_nonfinite, true, "Fail if the pre-clipping gradient norm is NaN or Inf."); DEFINE_string(clip_grad_foreach, "auto", "Gradient clipping path: auto|true|false."); // lr scheduler @@ -546,6 +546,14 @@ void Train(const nn::parallel::Rank &rank) { y = std::make_shared(y->To(device)); lossf = model->TrainStep({x}, {y}, optimizer, loss_fn, dtype); + if (optimizer->HasClipGradNormConfig()) { + auto *pp_model = dynamic_cast(model.get()); + auto norm_tensor = pp_model ? pp_model->last_grad_norm() : nullptr; + if (norm_tensor) { + auto total_grad_norm_cpu = norm_tensor->To(Device()); + total_grad_norm = *static_cast(total_grad_norm_cpu.DataPtr()); + } + } if (scheduler) { scheduler->Step(); } diff --git a/example/llama3/main.cc b/example/llama3/main.cc index c70cd9274..99a9e86f5 100644 --- a/example/llama3/main.cc +++ b/example/llama3/main.cc @@ -64,7 +64,7 @@ DEFINE_uint32(text_length, 64, "the length of the generated text"); DEFINE_double(learning_rate, 1e-5, "Peak learning rate."); DEFINE_int32(zero_stage, 0, "ZeRO stage (0/1/2/3); 0 disables DistributedOptimizer"); DEFINE_double(clip_grad_norm, -1.0, "Maximum gradient norm; negative disables clipping."); -DEFINE_double(grad_norm_type, 2.0, "Gradient norm type (positive finite p or inf)."); +DEFINE_double(grad_norm_type, 2.0, "Gradient norm type (finite p, inf, or -inf)."); DEFINE_bool(clip_grad_error_if_nonfinite, true, "Fail if the pre-clipping gradient norm is NaN or Inf."); DEFINE_string(clip_grad_foreach, "auto", "Gradient clipping path: auto|true|false."); // lr scheduler @@ -525,6 +525,14 @@ void Train(const nn::parallel::Rank &rank) { y = std::make_shared(y->To(device)); lossf = model->TrainStep({x}, {y}, optimizer, loss_fn, dtype); + if (optimizer->HasClipGradNormConfig()) { + auto *pp_model = dynamic_cast(model.get()); + auto norm_tensor = pp_model ? pp_model->last_grad_norm() : nullptr; + if (norm_tensor) { + auto total_grad_norm_cpu = norm_tensor->To(Device()); + total_grad_norm = *static_cast(total_grad_norm_cpu.DataPtr()); + } + } if (scheduler) { scheduler->Step(); } diff --git a/infini_train/include/nn/parallel/pp/pipeline_parallel.h b/infini_train/include/nn/parallel/pp/pipeline_parallel.h index 25939bdc2..a69f6babb 100644 --- a/infini_train/include/nn/parallel/pp/pipeline_parallel.h +++ b/infini_train/include/nn/parallel/pp/pipeline_parallel.h @@ -40,6 +40,8 @@ class PipelineParallel : public Module { std::vector> *mutable_chunks(); + std::shared_ptr last_grad_norm() const; + private: void BuildPipelineStage(const std::vector> &recv_shape, Device device, std::vector> &&chunks); diff --git a/infini_train/include/nn/parallel/pp/pipeline_schedule.h b/infini_train/include/nn/parallel/pp/pipeline_schedule.h index cae190f82..e621b73ca 100644 --- a/infini_train/include/nn/parallel/pp/pipeline_schedule.h +++ b/infini_train/include/nn/parallel/pp/pipeline_schedule.h @@ -31,12 +31,15 @@ class PipelineSchedule { const std::vector> &target_mbs, const std::shared_ptr &loss_fn, DataType dtype); + std::shared_ptr last_grad_norm() const { return last_grad_norm_; } + std::vector> ReceiveFromPrev(int peer_rank); std::vector> SendToNext(const std::vector> &tensors, int peer_rank); protected: int num_micro_batches_ = -1; std::shared_ptr stage_ = nullptr; + std::shared_ptr last_grad_norm_ = nullptr; }; class PipelineParallelScheduler { diff --git a/infini_train/include/optimizer.h b/infini_train/include/optimizer.h index 5b3007815..6b4cc8c50 100644 --- a/infini_train/include/optimizer.h +++ b/infini_train/include/optimizer.h @@ -41,7 +41,8 @@ class Optimizer { } // Scale selected gradients without replacing their storage. - void ScaleGradients_(const std::vector> ¶meters, float scale); + void ScaleGradients_(const std::vector> ¶meters, float scale, + bool foreach = false); void SetClipGradNormConfig(float max_norm, float norm_type = 2.0f, bool error_if_nonfinite = false, std::optional foreach = std::nullopt); diff --git a/infini_train/src/kernels/cpu/accumulate_grad.cc b/infini_train/src/kernels/cpu/accumulate_grad.cc index 25cf525e7..1a80b5f2d 100644 --- a/infini_train/src/kernels/cpu/accumulate_grad.cc +++ b/infini_train/src/kernels/cpu/accumulate_grad.cc @@ -1,5 +1,6 @@ #include #include +#include #include "infini_train/include/dispatcher.h" #include "infini_train/include/tensor.h" @@ -30,6 +31,53 @@ void ScaleInplace(const std::shared_ptr &tensor, float scale) { LOG(FATAL) << "ScaleInplace only supports floating point gradients."; } } + +template +void ScaleInplaceMultiTyped(const std::vector> &tensors, float scale) { + size_t total = 0; + for (const auto &tensor : tensors) { + if (tensor && tensor->NumElements() != 0) { + total += tensor->NumElements(); + } + } +#pragma omp parallel for + for (int64_t flat = 0; flat < static_cast(total); ++flat) { + size_t offset = static_cast(flat); + for (const auto &tensor : tensors) { + if (!tensor || tensor->NumElements() == 0 || offset >= tensor->NumElements()) { + if (tensor && tensor->NumElements() != 0) { + offset -= tensor->NumElements(); + } + continue; + } + auto *data = static_cast(tensor->DataPtr()); + data[offset] = T(static_cast(data[offset]) * scale); + break; + } + } +} + +void ScaleInplaceMulti(std::vector> tensors, float scale) { + if (tensors.empty() || scale == 1.0f) { + return; + } + switch (tensors.front()->Dtype()) { + case DataType::kFLOAT16: + ScaleInplaceMultiTyped(tensors, scale); + return; + case DataType::kBFLOAT16: + ScaleInplaceMultiTyped(tensors, scale); + return; + case DataType::kFLOAT32: + ScaleInplaceMultiTyped(tensors, scale); + return; + case DataType::kFLOAT64: + ScaleInplaceMultiTyped(tensors, scale); + return; + default: + LOG(FATAL) << "ScaleInplaceMulti only supports floating point gradients."; + } +} void AccumulateGrad(const std::shared_ptr &gradient, float rate, const std::shared_ptr &tensor) { for (int64_t idx = 0; idx < gradient->NumElements(); ++idx) { static_cast(tensor->DataPtr())[idx] += rate * static_cast(gradient->DataPtr())[idx]; @@ -66,6 +114,7 @@ void AdamAccumulateGrad(const std::shared_ptr &grad, const std::shared_p REGISTER_CPU_ACCUMULATE_GRAD_KERNEL(AccumulateGrad) REGISTER_CPU_ACCUMULATE_GRAD_KERNEL(ScaleInplace) +REGISTER_CPU_ACCUMULATE_GRAD_KERNEL(ScaleInplaceMulti) REGISTER_CPU_ACCUMULATE_GRAD_KERNEL(AdamAccumulateGrad) #undef REGISTER_CPU_ACCUMULATE_GRAD_KERNEL diff --git a/infini_train/src/kernels/cuda/accumulate_grad.cu b/infini_train/src/kernels/cuda/accumulate_grad.cu index 808a51b89..64ed9bc9e 100644 --- a/infini_train/src/kernels/cuda/accumulate_grad.cu +++ b/infini_train/src/kernels/cuda/accumulate_grad.cu @@ -1,6 +1,8 @@ #include #include +#include +#include "infini_train/include/common/cuda/common_cuda.h" #include "infini_train/include/common/cuda/kernel_helper.cuh" #include "infini_train/include/core/runtime/device_guard.h" #include "infini_train/include/dispatcher.h" @@ -35,6 +37,73 @@ void ScaleInplace(const std::shared_ptr &tensor, float scale) { "CUDA ScaleInplace"); } +template +__global__ void ScaleInplaceMultiKernel(T **ptrs, const size_t *offsets, size_t total, int num_tensors, float scale) { + const size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total) { + return; + } + int lo = 0; + int hi = num_tensors; + while (lo + 1 < hi) { + const int mid = lo + (hi - lo) / 2; + if (offsets[mid] <= idx) { + lo = mid; + } else { + hi = mid; + } + } + const size_t local_idx = idx - offsets[lo]; + ptrs[lo][local_idx] = common::cuda::Mul(ptrs[lo][local_idx], common::cuda::Cast(scale)); +} + +void ScaleInplaceMulti(std::vector> tensors, float scale) { + if (tensors.empty() || scale == 1.0f) { + return; + } + std::vector host_ptrs; + std::vector host_offsets; + host_ptrs.reserve(tensors.size()); + host_offsets.reserve(tensors.size() + 1); + host_offsets.push_back(0); + for (const auto &tensor : tensors) { + if (!tensor || tensor->NumElements() == 0) { + continue; + } + host_ptrs.push_back(tensor->DataPtr()); + host_offsets.push_back(host_offsets.back() + tensor->NumElements()); + } + const size_t total = host_offsets.back(); + if (total == 0 || host_ptrs.empty()) { + return; + } + + auto device = tensors.front()->GetDevice(); + const auto &cuda_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); + void *device_ptrs = nullptr; + size_t *device_offsets = nullptr; + CUDA_CHECK(cudaMallocAsync(&device_ptrs, sizeof(void *) * host_ptrs.size(), cuda_stream)); + CUDA_CHECK(cudaMallocAsync(&device_offsets, sizeof(size_t) * host_offsets.size(), cuda_stream)); + CUDA_CHECK(cudaMemcpyAsync(device_ptrs, host_ptrs.data(), sizeof(void *) * host_ptrs.size(), + cudaMemcpyHostToDevice, cuda_stream)); + CUDA_CHECK(cudaMemcpyAsync(device_offsets, host_offsets.data(), sizeof(size_t) * host_offsets.size(), + cudaMemcpyHostToDevice, cuda_stream)); + + const int threads_per_block = 256; + const int num_blocks = static_cast((total + threads_per_block - 1) / threads_per_block); + core::cuda::DispatchCudaFunc( + tensors.front()->Dtype(), + [=]() { + ScaleInplaceMultiKernel<<>>( + static_cast(device_ptrs), device_offsets, total, static_cast(host_ptrs.size()), scale); + }, + "CUDA ScaleInplaceMulti"); + CUDA_CHECK(cudaFreeAsync(device_ptrs, cuda_stream)); + CUDA_CHECK(cudaFreeAsync(device_offsets, cuda_stream)); +} + template __global__ void AccumulateGradKernel(const T *grad_ptr, float rate, T *tensor_ptr, size_t num_elements) { int idx = blockIdx.x * blockDim.x + threadIdx.x; @@ -115,6 +184,7 @@ void AdamAccumulateGrad(const std::shared_ptr &grad, const std::shared_p REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL(AccumulateGrad) REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL(ScaleInplace) +REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL(ScaleInplaceMulti) REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL(AdamAccumulateGrad) #undef REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL diff --git a/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc b/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc index 11dda6e7b..7c648a202 100644 --- a/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc +++ b/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc @@ -165,11 +165,8 @@ std::shared_ptr DistributedOptimizer::ClipGradNorm_(const std::vector foreach) { CHECK_GE(max_norm, 0.0f) << "max_norm must be non-negative."; - CHECK((norm_type > 0.0f && std::isfinite(norm_type)) || norm_type == std::numeric_limits::infinity()) - << "norm_type must be positive finite or +inf."; - if (foreach.value_or(false)) { - LOG(FATAL) << "DistributedOptimizer: foreach=true is not implemented."; - } + CHECK(std::isfinite(norm_type) || std::isinf(norm_type)) + << "norm_type must be finite, +inf, or -inf."; FinishGradSync(); @@ -204,8 +201,13 @@ std::shared_ptr DistributedOptimizer::ClipGradNorm_(const std::vectorClipGradNorm_(norm_shards, std::numeric_limits::infinity(), norm_type, false, std::nullopt); const float local_norm = *static_cast(local_norm_tensor->DataPtr()); - const bool is_inf_norm = std::isinf(norm_type); - double local_stat = is_inf_norm ? static_cast(local_norm) + const bool is_pos_inf_norm = norm_type == std::numeric_limits::infinity(); + const bool is_neg_inf_norm = norm_type == -std::numeric_limits::infinity(); + const bool is_zero_norm = norm_type == 0.0f; + double local_stat = is_neg_inf_norm && norm_shards.empty() + ? std::numeric_limits::infinity() + : is_pos_inf_norm || is_neg_inf_norm ? static_cast(local_norm) + : is_zero_norm ? static_cast(local_norm) : std::pow(static_cast(local_norm), static_cast(norm_type)); const ProcessGroup *group = nullptr; @@ -227,30 +229,40 @@ std::shared_ptr DistributedOptimizer::ClipGradNorm_(const std::vector(std::vector{}, DataType::kFLOAT32, total_norm_device); reduced->Fill(static_cast(local_stat)); if (group && ddp_world_size_ > 1) { - group->AllReduce(reduced, is_inf_norm ? function::ReduceOpType::kMax : function::ReduceOpType::kSum, false); + group->AllReduce(reduced, is_neg_inf_norm ? function::ReduceOpType::kMin + : is_pos_inf_norm ? function::ReduceOpType::kMax + : function::ReduceOpType::kSum, + false); } if (global::GetTensorParallelSize() > 1) { const auto *tp_group = ProcessGroupFactory::Instance(total_norm_device.type()) ->Get(GetTensorParallelProcessGroupName(total_norm_device.Rank().GlobalRank())); CHECK(tp_group) << "Tensor-parallel process group is not initialized."; - tp_group->AllReduce(reduced, is_inf_norm ? function::ReduceOpType::kMax : function::ReduceOpType::kSum, false); + tp_group->AllReduce(reduced, is_neg_inf_norm ? function::ReduceOpType::kMin + : is_pos_inf_norm ? function::ReduceOpType::kMax + : function::ReduceOpType::kSum, + false); } if (nn::parallel::global::GetPipelineParallelSize() > 1) { const auto *pp_group = ProcessGroupFactory::Instance(total_norm_device.type()) ->Get(GetPipelineParallelProcessGroupName(total_norm_device.Rank().GlobalRank())); CHECK(pp_group) << "Pipeline process group is not initialized."; - pp_group->AllReduce(reduced, is_inf_norm ? function::ReduceOpType::kMax : function::ReduceOpType::kSum, false); + pp_group->AllReduce(reduced, is_neg_inf_norm ? function::ReduceOpType::kMin + : is_pos_inf_norm ? function::ReduceOpType::kMax + : function::ReduceOpType::kSum, + false); } Tensor reduced_cpu = reduced->GetDevice().IsCPU() ? Tensor(*reduced, 0, reduced->Dims()) : reduced->To(Device()); const float reduced_value = *static_cast(reduced_cpu.DataPtr()); - const double total_norm = is_inf_norm ? static_cast(reduced_value) - : std::pow(static_cast(reduced_value), 1.0 / norm_type); + const double total_norm = is_zero_norm || is_pos_inf_norm || is_neg_inf_norm + ? static_cast(reduced_value) + : std::pow(static_cast(reduced_value), 1.0 / norm_type); if (error_if_nonfinite && !std::isfinite(total_norm)) { LOG(FATAL) << "The total gradient norm is non-finite."; } const double coefficient = std::min(static_cast(max_norm) / (total_norm + 1e-6), 1.0); - base_optimizer_->ScaleGradients_(selected_shards, static_cast(coefficient)); + base_optimizer_->ScaleGradients_(selected_shards, static_cast(coefficient), foreach.value_or(false)); auto result = std::make_shared(std::vector{}, DataType::kFLOAT32, Device()); *static_cast(result->DataPtr()) = static_cast(total_norm); diff --git a/infini_train/src/nn/parallel/pp/pipeline_parallel.cc b/infini_train/src/nn/parallel/pp/pipeline_parallel.cc index c0369cdeb..4eee44d8c 100644 --- a/infini_train/src/nn/parallel/pp/pipeline_parallel.cc +++ b/infini_train/src/nn/parallel/pp/pipeline_parallel.cc @@ -104,4 +104,8 @@ PipelineParallel::PipelineParallel(const std::shared_ptr module, int num } std::vector> *PipelineParallel::mutable_chunks() { return pipeline_stage_->mutable_chunks(); } + +std::shared_ptr PipelineParallel::last_grad_norm() const { + return schedule_ ? schedule_->last_grad_norm() : nullptr; +} } // namespace infini_train::nn::parallel diff --git a/infini_train/src/nn/parallel/pp/pipeline_schedule.cc b/infini_train/src/nn/parallel/pp/pipeline_schedule.cc index 7a507d1b8..5b9f1cdc7 100644 --- a/infini_train/src/nn/parallel/pp/pipeline_schedule.cc +++ b/infini_train/src/nn/parallel/pp/pipeline_schedule.cc @@ -300,7 +300,7 @@ float PipelineSchedule::Step(std::shared_ptr input, std::shared_ptrClipGradNormConfigured(); + last_grad_norm_ = optimizer->ClipGradNormConfigured(); optimizer->Step(); return lossf; diff --git a/infini_train/src/optimizer.cc b/infini_train/src/optimizer.cc index 26d922263..6d2f490e5 100644 --- a/infini_train/src/optimizer.cc +++ b/infini_train/src/optimizer.cc @@ -3,6 +3,8 @@ #include #include #include +#include +#include #include #include #include @@ -40,6 +42,8 @@ namespace { struct GradientNormStats { double sum = 0.0; double max_abs = 0.0; + double min_abs = std::numeric_limits::infinity(); + uint64_t nonzero = 0; bool has_gradient = false; bool finite = true; }; @@ -51,29 +55,54 @@ GradientNormStats ComputeGradientNormStats(const std::vectorgrad()) { continue; } - stats.has_gradient = true; const auto &gradient = parameter->grad(); - Tensor host = gradient->GetDevice().IsCPU() ? Tensor(*gradient, 0, gradient->Dims()) : gradient->To(Device()); + if (gradient->NumElements() == 0) { + continue; + } + stats.has_gradient = true; + std::shared_ptr host; + if (gradient->GetDevice().IsCPU()) { + host = std::make_shared(*gradient, 0, gradient->Dims()); + } else { + host = std::make_shared(gradient->To(Device())); + } if (gradient->GetDevice().IsCUDA()) { core::GetDeviceGuardImpl(gradient->GetDevice().type())->SynchronizeDevice(gradient->GetDevice()); } - Tensor fp32 = host.Dtype() == DataType::kFLOAT32 ? host : host.To(DataType::kFLOAT32); - const float *data = static_cast(fp32.DataPtr()); - for (size_t i = 0; i < fp32.NumElements(); ++i) { + std::shared_ptr fp32; + if (host->Dtype() == DataType::kFLOAT32) { + fp32 = host; + } else { + fp32 = std::make_shared(host->To(DataType::kFLOAT32)); + } + const float *data = static_cast(fp32->DataPtr()); + bool tensor_has_nonzero = false; + for (size_t i = 0; i < fp32->NumElements(); ++i) { const double value = static_cast(data[i]); const double abs_value = std::abs(value); if (!std::isfinite(value)) { stats.finite = false; - if (std::isinf(norm_type)) { + if (norm_type == -std::numeric_limits::infinity()) { + stats.min_abs = std::numeric_limits::quiet_NaN(); + } else if (std::isinf(norm_type)) { stats.max_abs = std::numeric_limits::quiet_NaN(); } } - if (std::isinf(norm_type)) { + if (norm_type == -std::numeric_limits::infinity()) { + stats.min_abs = std::min(stats.min_abs, abs_value); + } else if (std::isinf(norm_type)) { stats.max_abs = std::max(stats.max_abs, abs_value); + } else if (norm_type == 0.0f) { + if (abs_value != 0.0) { + tensor_has_nonzero = true; + } } else { stats.sum += std::pow(abs_value, static_cast(norm_type)); } } + if (norm_type == 0.0f && tensor_has_nonzero) { + ++stats.nonzero; + } } return stats; } @@ -108,13 +137,32 @@ std::shared_ptr Optimizer::ClipGradNormConfigured() { return ClipGradNorm_(params_, config.max_norm, config.norm_type, config.error_if_nonfinite, config.foreach); } -void Optimizer::ScaleGradients_(const std::vector> ¶meters, float scale) { +void Optimizer::ScaleGradients_(const std::vector> ¶meters, float scale, bool foreach) { std::unordered_set seen; + using BatchKey = std::tuple; + std::map>> batches; for (const auto ¶meter : parameters) { if (!parameter || !seen.insert(parameter.get()).second || !parameter->grad()) { continue; } - ScaleGradientInplace(parameter->grad(), scale); + if (foreach) { + const auto device = parameter->grad()->GetDevice(); + batches[{static_cast(device.type()), static_cast(device.index()), + static_cast(parameter->grad()->Dtype())}] + .push_back(parameter->grad()); + } else { + ScaleGradientInplace(parameter->grad(), scale); + } + } + if (foreach) { + for (const auto &[_, batch] : batches) { + if (batch.empty()) { + continue; + } + core::DeviceGuard guard(batch.front()->GetDevice()); + auto kernel = Dispatcher::Instance().GetKernel({batch.front()->GetDevice().type(), "ScaleInplaceMulti"}); + kernel.Call(batch, scale); + } } } @@ -122,16 +170,21 @@ std::shared_ptr Optimizer::ClipGradNorm_(const std::vector foreach) { CHECK_GE(max_norm, 0.0f) << "max_norm must be non-negative."; - CHECK((norm_type > 0.0f && std::isfinite(norm_type)) || norm_type == std::numeric_limits::infinity()) - << "norm_type must be positive finite or +inf."; - if (foreach.value_or(false)) { - LOG(FATAL) << "ClipGradNorm_: foreach=true is not implemented."; - } + CHECK(std::isfinite(norm_type) || std::isinf(norm_type)) + << "norm_type must be finite, +inf, or -inf."; const auto stats = ComputeGradientNormStats(parameters, norm_type); double total_norm = 0.0; if (stats.has_gradient) { - total_norm = std::isinf(norm_type) ? stats.max_abs : std::pow(stats.sum, 1.0 / norm_type); + if (norm_type == 0.0f) { + total_norm = static_cast(stats.nonzero); + } else if (norm_type == -std::numeric_limits::infinity()) { + total_norm = stats.min_abs; + } else if (norm_type == std::numeric_limits::infinity()) { + total_norm = stats.max_abs; + } else { + total_norm = std::pow(stats.sum, 1.0 / norm_type); + } } // Pipeline stages own disjoint parameter sets. Reduce their statistics so @@ -146,27 +199,39 @@ std::shared_ptr Optimizer::ClipGradNorm_(const std::vector(std::vector{}, DataType::kFLOAT32, device); - reduced->Fill(static_cast(std::isinf(norm_type) ? stats.max_abs : stats.sum)); + const double local_stat = norm_type == 0.0f ? static_cast(stats.nonzero) + : norm_type == -std::numeric_limits::infinity() ? stats.min_abs + : norm_type == std::numeric_limits::infinity() ? stats.max_abs + : stats.sum; + reduced->Fill(static_cast(local_stat)); const auto *group = nn::parallel::ProcessGroupFactory::Instance(device.type()) ->Get(nn::parallel::GetPipelineParallelProcessGroupName(device.Rank().GlobalRank())); CHECK(group) << "Pipeline process group is not initialized."; group->AllReduce(reduced, - std::isinf(norm_type) ? nn::parallel::function::ReduceOpType::kMax - : nn::parallel::function::ReduceOpType::kSum, + norm_type == -std::numeric_limits::infinity() + ? nn::parallel::function::ReduceOpType::kMin + : norm_type == std::numeric_limits::infinity() + ? nn::parallel::function::ReduceOpType::kMax + : nn::parallel::function::ReduceOpType::kSum, false); Tensor reduced_cpu = reduced->GetDevice().IsCPU() ? Tensor(*reduced, 0, reduced->Dims()) : reduced->To(Device()); const float value = *static_cast(reduced_cpu.DataPtr()); - total_norm = std::isinf(norm_type) ? value : std::pow(static_cast(value), 1.0 / norm_type); + total_norm = norm_type == 0.0f ? static_cast(value) + : norm_type == -std::numeric_limits::infinity() + ? static_cast(value) + : norm_type == std::numeric_limits::infinity() + ? static_cast(value) + : std::pow(static_cast(value), 1.0 / norm_type); } - if (error_if_nonfinite && (!std::isfinite(total_norm))) { + if (error_if_nonfinite && (!stats.finite || !std::isfinite(total_norm))) { LOG(FATAL) << "The total gradient norm is non-finite."; } const double coefficient = std::isinf(max_norm) ? 1.0 : std::min(static_cast(max_norm) / (total_norm + 1e-6), 1.0); if (stats.has_gradient) { - ScaleGradients_(parameters, static_cast(coefficient)); + ScaleGradients_(parameters, static_cast(coefficient), foreach.value_or(false)); } return MakeScalar(static_cast(total_norm)); } diff --git a/scripts/test_config.json b/scripts/test_config.json index 2b8202d27..0bde46426 100644 --- a/scripts/test_config.json +++ b/scripts/test_config.json @@ -15,8 +15,8 @@ "RUN_PROFILE_TEST": "true", "MIXTRAL_INPUT_BIN": "/data1/shared/InfiniTrain-dev/data/llmc/llama3/tinyshakespeare/tiny_shakespeare_train.bin", "MIXTRAL_LLMC_FILEPATH": "/data1/shared/InfiniTrain-dev/data/llmc/mixtral/mixtral_megatron_export.bin", - "GPT2_TEST_GROUPS": "basic,zero,lora,checkpoint,lr_scheduler,8_proc", - "LLAMA3_TEST_GROUPS": "basic,zero,lora,checkpoint,lr_scheduler,8_proc", + "GPT2_TEST_GROUPS": "basic,zero,lora,checkpoint,lr_scheduler,8_proc,gradient_clipping_excellent", + "LLAMA3_TEST_GROUPS": "basic,zero,lora,checkpoint,lr_scheduler,8_proc,gradient_clipping_excellent", "MIXTRAL_TEST_GROUPS": "moe" }, "basic_compile_commands": [ @@ -26,6 +26,72 @@ } ], "test_groups": [ + { + "tag": "gradient_clipping_excellent", + "tests": [ + { + "id": "single_fp32_foreach", + "args": { + "dtype": "float32", + "num_iteration": 1, + "batch_size": 1, + "total_batch_size": 64, + "sequence_length": 64, + "overfit_single_batch": true, + "clip_grad_norm": 0.1, + "grad_norm_type": 2.0, + "clip_grad_error_if_nonfinite": true, + "clip_grad_foreach": "true" + } + }, + { + "id": "single_bf16_p0", + "args": { + "dtype": "bfloat16", + "num_iteration": 1, + "batch_size": 1, + "total_batch_size": 64, + "sequence_length": 64, + "overfit_single_batch": true, + "clip_grad_norm": 0.1, + "grad_norm_type": 0.0, + "clip_grad_foreach": "auto" + } + }, + { + "id": "tp4_sp_zero2", + "args": { + "dtype": "float32", + "num_iteration": 1, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 4, + "sequence_parallel": true, + "zero_stage": 2, + "clip_grad_norm": 1.0, + "grad_norm_type": 2.0, + "clip_grad_foreach": "true" + } + }, + { + "id": "pp2_vpp2_lora", + "args": { + "dtype": "float32", + "num_iteration": 1, + "batch_size": 40, + "total_batch_size": 5120, + "pipeline_parallel": 2, + "virtual_pipeline_parallel": 2, + "lora_rank": 4, + "lora_alpha": 8.0, + "lora_target_modules": "c_fc,c_proj", + "clip_grad_norm": 1.0, + "grad_norm_type": 2.0, + "clip_grad_foreach": "auto" + } + } + ] + }, { "tag": "basic", "tests": [ diff --git a/tests/optimizer/test_clip_grad_norm.cc b/tests/optimizer/test_clip_grad_norm.cc index 92e4ffa7c..69fd39622 100644 --- a/tests/optimizer/test_clip_grad_norm.cc +++ b/tests/optimizer/test_clip_grad_norm.cc @@ -28,6 +28,15 @@ static std::shared_ptr MakeTensor(Device device, const std::vector MakeTensor(Device device, DataType dtype, const std::vector &values) { + auto fp32 = std::make_shared( + values.data(), std::vector{static_cast(values.size())}, DataType::kFLOAT32, device); + auto converted = fp32->To(dtype); + auto result = std::make_shared(converted.Dims(), dtype, device); + result->CopyFrom(converted); + return result; +} + static float ScalarCPU(const std::shared_ptr &value) { auto cpu = value->To(Device()); return *static_cast(cpu.DataPtr()); @@ -103,6 +112,57 @@ TEST_P(ClipGradNormTest, NonIntegerPAndZeroMaxNorm) { EXPECT_FLOAT_EQ(values[1], 0.0f); } +TEST_P(ClipGradNormTest, SupportsZeroNegativeAndNegativeInfinityNorms) { + auto param = std::make_shared(std::vector{3}, DataType::kFLOAT32, GetDevice()); + auto grad = MakeTensor(GetDevice(), {-2.0f, 0.0f, 4.0f}); + param->set_grad(grad); + auto optimizer = std::make_shared(std::vector>{param}, 0.1f); + + auto zero = optimizer->ClipGradNorm({param}, 1.0f, 0.0f); + // PyTorch computes the 0-norm of each gradient tensor and then the + // 0-norm of that list, so one non-empty gradient tensor contributes one. + EXPECT_FLOAT_EQ(ScalarCPU(zero), 1.0f); + + auto neg_grad = MakeTensor(GetDevice(), {-2.0f, 0.0f, 4.0f}); + param->set_grad(neg_grad); + auto negative = optimizer->ClipGradNorm({param}, 1.0f, -1.0f); + EXPECT_TRUE(std::isfinite(ScalarCPU(negative)) || ScalarCPU(negative) == 0.0f); + + auto min_grad = MakeTensor(GetDevice(), {-2.0f, 0.5f, 4.0f}); + param->set_grad(min_grad); + auto neg_inf = optimizer->ClipGradNorm({param}, 0.25f, -std::numeric_limits::infinity()); + EXPECT_NEAR(ScalarCPU(neg_inf), 0.5f, 1e-5f); +} + +TEST_P(ClipGradNormTest, ForeachTrueUsesMultiTensorSemantics) { + auto first = std::make_shared(std::vector{2}, DataType::kFLOAT32, GetDevice()); + auto second = std::make_shared(std::vector{2}, DataType::kFLOAT32, GetDevice()); + first->set_grad(MakeTensor(GetDevice(), {3.0f, 4.0f})); + second->set_grad(MakeTensor(GetDevice(), {0.0f, 12.0f})); + auto optimizer = std::make_shared( + std::vector>{first, second}, 0.1f); + auto total = optimizer->ClipGradNorm({first, second}, 6.5f, 2.0f, false, true); + EXPECT_NEAR(ScalarCPU(total), 13.0f, 1e-5f); + auto clipped = second->grad()->To(Device()); + EXPECT_NEAR(static_cast(clipped.DataPtr())[1], 6.0f, 1e-4f); +} + +TEST_P(ClipGradNormTest, AccumulatesReducedPrecisionGradientsInFloat32) { + auto fp16_param = std::make_shared(std::vector{2}, DataType::kFLOAT16, GetDevice()); + auto bf16_param = std::make_shared(std::vector{2}, DataType::kBFLOAT16, GetDevice()); + fp16_param->set_grad(MakeTensor(GetDevice(), DataType::kFLOAT16, {3.0f, 4.0f})); + bf16_param->set_grad(MakeTensor(GetDevice(), DataType::kBFLOAT16, {0.0f, 12.0f})); + auto optimizer = std::make_shared( + std::vector>{fp16_param, bf16_param}, 0.1f); + + auto total = optimizer->ClipGradNorm({fp16_param, bf16_param}, 6.5f, 2.0f, false, true); + EXPECT_NEAR(ScalarCPU(total), 13.0f, 2e-2f); + auto fp16_cpu = fp16_param->grad()->To(Device()); + auto bf16_cpu = bf16_param->grad()->To(Device()); + EXPECT_NEAR(static_cast(static_cast(fp16_cpu.DataPtr())[0]), 1.5f, 8e-2f); + EXPECT_NEAR(static_cast(static_cast(bf16_cpu.DataPtr())[1]), 6.0f, 1.5e-1f); +} + TEST_P(ClipGradNormTest, ErrorOnNonFiniteBeforeScaling) { ONLY_CPU(); auto param = std::make_shared(std::vector{2}, DataType::kFLOAT32, GetDevice()); @@ -110,4 +170,17 @@ TEST_P(ClipGradNormTest, ErrorOnNonFiniteBeforeScaling) { param->set_grad(grad); auto optimizer = std::make_shared(std::vector>{param}, 0.1f); EXPECT_DEATH(optimizer->ClipGradNorm({param}, 1.0f, 2.0f, true), "non-finite"); + + auto neg_inf_grad = MakeTensor(GetDevice(), {std::numeric_limits::quiet_NaN(), 1.0f}); + param->set_grad(neg_inf_grad); + EXPECT_DEATH(optimizer->ClipGradNorm({param}, 1.0f, -std::numeric_limits::infinity(), true), "non-finite"); +} + +TEST_P(ClipGradNormTest, EmptyGradientIsIgnoredForNegativeInfinity) { + auto param = std::make_shared(std::vector{0}, DataType::kFLOAT32, GetDevice()); + auto grad = std::make_shared(std::vector{0}, DataType::kFLOAT32, GetDevice()); + param->set_grad(grad); + auto optimizer = std::make_shared(std::vector>{param}, 0.1f); + auto total = optimizer->ClipGradNorm({param}, 1.0f, -std::numeric_limits::infinity()); + EXPECT_FLOAT_EQ(ScalarCPU(total), 0.0f); } diff --git a/tests/optimizer/test_optimizer_parameter_names.cc b/tests/optimizer/test_optimizer_parameter_names.cc index c96f5ca4b..32c33324a 100644 --- a/tests/optimizer/test_optimizer_parameter_names.cc +++ b/tests/optimizer/test_optimizer_parameter_names.cc @@ -1,3 +1,4 @@ +#include #include #include @@ -123,7 +124,7 @@ TEST_P(OptimizerParameterNamesTest, DistributedOptimizerClipGradNormUsesZero2Loc full_grad->Fill(1.0f); auto group = ddp_model->bucket_groups().front(); group->AccumulateParamGrad(parameter, full_grad, /*overwrite=*/true, /*learning_rate=*/1.0f); - auto total_norm = optimizer.ClipGradNorm(model->Parameters(), 8.0f, 2.0f); + auto total_norm = optimizer.ClipGradNorm(model->Parameters(), 8.0f, 2.0f, false, true); auto total_norm_cpu = total_norm->To(Device()); EXPECT_NEAR(*static_cast(total_norm_cpu.DataPtr()), 16.0f, 1e-4f); auto local_shard = group->GetLocalGradShardBuffer(0)->To(Device()); @@ -154,7 +155,7 @@ TEST_P(OptimizerParameterNamesTest, DistributedOptimizerClipGradNormHandlesEmpty auto parameter = model->Parameters().front(); ASSERT_NE(parameter->grad(), nullptr); parameter->grad()->Fill(1.0f); - auto total_norm = optimizer.ClipGradNorm(model->Parameters(), 2.0f, 2.0f); + auto total_norm = optimizer.ClipGradNorm(model->Parameters(), 2.0f, 2.0f, false, true); auto total_norm_cpu = total_norm->To(Device()); EXPECT_NEAR(*static_cast(total_norm_cpu.DataPtr()), 4.0f, 1e-4f); } @@ -191,7 +192,46 @@ TEST_P(OptimizerParameterNamesTest, DistributedOptimizerClipGradNormUsesGlobalSh parameter->grad()->Fill(1.0f); } - auto total_norm = optimizer.ClipGradNorm(model->Parameters(), 2.0f, 2.0f); + auto total_norm = optimizer.ClipGradNorm(model->Parameters(), 2.0f, 2.0f, false, true); auto total_norm_cpu = total_norm->To(Device()); EXPECT_NEAR(*static_cast(total_norm_cpu.DataPtr()), 16.0f, 1e-4f); } + +TEST_P(OptimizerParameterNamesTest, DistributedOptimizerClipGradNormSupportsZeroAndNegativeInfinity) { + ONLY_CUDA(); + REQUIRE_MIN_DEVICES(2); + if (nn::parallel::global::GetDataParallelSize() != 2) { + GTEST_SKIP() << "requires WORLD_SIZE=2"; + } + const nn::parallel::Rank rank(nn::parallel::global::GetGlobalProcRank(), /*thread_rank=*/0, + nn::parallel::global::GetNprocPerNode(), /*thread_size=*/1); + auto *pg_factory = nn::parallel::ProcessGroupFactory::Instance(Device::DeviceType::kCUDA); + pg_factory->GetOrCreate(nn::parallel::GetDataParallelProcessGroupName(rank.GlobalRank()), + nn::parallel::GetDataParallelGroupRanks(rank.GlobalRank())); + auto model = std::make_shared( + 64, 4, /*bias=*/false, Device(Device::DeviceType::kCUDA, nn::parallel::global::GetLocalProcRank())); + nn::parallel::DistributedDataParallelConfig ddp_config; + ddp_config.zero_stage = 1; + ddp_config.overlap_grad_reduce = false; + ddp_config.overlap_param_gather = false; + auto ddp_model = std::make_shared(model, rank, ddp_config); + nn::parallel::DistributedOptimizer optimizer(optimizers::SGD::CreateNamed(0.1f), model->NamedParameters(), + std::vector>{ddp_model}, + /*ddp_world_size=*/2, /*ddp_rank=*/rank.GlobalRank()); + for (const auto ¶meter : model->Parameters()) { + ASSERT_NE(parameter->grad(), nullptr); + parameter->grad()->Fill(1.0f); + } + + auto zero_norm = optimizer.ClipGradNorm(model->Parameters(), 1.0f, 0.0f, false, true); + auto zero_cpu = zero_norm->To(Device()); + EXPECT_NEAR(*static_cast(zero_cpu.DataPtr()), 2.0f, 1e-4f); + + for (const auto ¶meter : model->Parameters()) { + parameter->grad()->Fill(1.0f); + } + auto neg_inf_norm = + optimizer.ClipGradNorm(model->Parameters(), 0.5f, -std::numeric_limits::infinity(), false, true); + auto neg_inf_cpu = neg_inf_norm->To(Device()); + EXPECT_NEAR(*static_cast(neg_inf_cpu.DataPtr()), 1.0f, 1e-4f); +} From 33742085869e41930ded7826ea55634f2c1c7020 Mon Sep 17 00:00:00 2001 From: zhengjiahong Date: Tue, 15 Sep 2026 20:38:40 +0800 Subject: [PATCH 7/7] =?UTF-8?q?=E6=9B=B4=E6=96=B0=E6=8A=A5=E5=91=8A?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/gradient_norm_clipping.md | 120 +++++++++-- gradient_norm_clipping_acceptance_report.md | 214 ++++++++++++++++++++ 2 files changed, 321 insertions(+), 13 deletions(-) create mode 100644 gradient_norm_clipping_acceptance_report.md diff --git a/docs/gradient_norm_clipping.md b/docs/gradient_norm_clipping.md index 41ed90d0b..033b79865 100644 --- a/docs/gradient_norm_clipping.md +++ b/docs/gradient_norm_clipping.md @@ -1,22 +1,116 @@ -# Distributed gradient norm and clipping +# 分布式梯度范数计算与梯度裁剪 -The public entry point is: +## 公开接口 - auto total_norm = optimizer->ClipGradNorm_(parameters, max_norm, norm_type, error_if_nonfinite, std::nullopt); +公开调用入口如下： -It returns a CPU FP32 scalar containing the pre-clipping norm and scales the selected gradients in place. Parameters without gradients are ignored and duplicate parameter pointers are counted once. `norm_type` follows PyTorch-style vector norms: `0` first counts each gradient tensor that contains a non-zero element and then takes the 0-norm of that tensor list, finite positive and negative p use the p-norm formula, `+inf` takes the maximum absolute value, and `-inf` takes the minimum absolute value. The coefficient is min(max_norm / (total_norm + 1e-6), 1); max_norm == 0 therefore zeros gradients rather than disabling clipping. +```cpp +auto total_norm = optimizer->ClipGradNorm_( + parameters, + max_norm, + norm_type, + error_if_nonfinite, + std::nullopt); +``` -`foreach=true` selects the multi-tensor scaling path. CPU and CUDA `ScaleInplaceMulti` kernels write through the existing gradient views, preserving flat-buffer and ZeRO-shard aliases while issuing one batch dispatch per device/dtype group. CUDA norm inspection synchronizes only to obtain scalar statistics; scaling remains on the CUDA stream. +该接口返回一个位于 CPU 上的 FP32 标量，表示**裁剪前的梯度总范数**，并对选中的梯度进行原地缩放。 -DistributedOptimizer::ClipGradNorm_ first finishes pending gradient collectives, computes the norm over the shard parameters consumed by its base optimizer, reduces statistics over the DP process group, and applies the same coefficient to every local shard. When PP is enabled, statistics are also reduced over the PP process group. Pipeline scheduling calls ClipGradNormConfigured after all micro-batches complete. +- 没有梯度的参数会被忽略。 +- 对于重复出现的参数指针，仅计算一次。 -The GPT-2 and LLaMA-3 examples expose: +## 梯度范数类型 -- --clip_grad_norm=-1 (negative disables clipping) -- --grad_norm_type=2 -- --clip_grad_error_if_nonfinite=true -- --clip_grad_foreach=auto|true|false +`norm_type` 遵循 PyTorch 风格的向量范数定义： -A total_grad_norm log line is emitted when clipping is enabled. +- `0`：首先判断每个梯度张量是否至少包含一个非零元素，然后对这些梯度张量组成的列表计算 0-范数，即统计包含非零元素的梯度张量数量。 +- 有限的正数或负数 `p`：按照 p-范数公式进行计算。 +- `+inf`：取所有梯度元素绝对值的最大值。 +- `-inf`：取所有梯度元素绝对值的最小值。 -Tensor-parallel replicated-parameter ownership is filtered by DistributedOptimizer; complete mixed TP/SP/PP/vPP training coverage is provided by the optional matrix in `scripts/test_config.json` and should be run on a host with the corresponding checkpoints. +梯度缩放系数的计算公式为： + +```text +min(max_norm / (total_norm + 1e-6), 1) +``` + +因此，当 `max_norm == 0` 时，所有梯度都会被置零，而不是禁用梯度裁剪。 + +## 多张量缩放 + +当设置 `foreach=true` 时，将使用多张量批量缩放路径。 + +CPU 和 CUDA 的 `ScaleInplaceMulti` 内核会直接写入现有的梯度视图，从而保留其与以下存储区域之间的别名关系： + +- 扁平缓冲区（flat buffer） +- ZeRO 梯度分片 + +同时，每个设备和数据类型分组仅进行一次批量分派。 + +CUDA 梯度范数检查仅在获取标量统计信息时执行同步；梯度缩放操作仍在 CUDA 流上执行。 + +## 分布式梯度裁剪 + +`DistributedOptimizer::ClipGradNorm_` 的处理流程如下： + +1. 等待所有尚未完成的梯度集合通信操作结束。 +2. 针对底层优化器实际使用的分片参数计算梯度范数。 +3. 在数据并行（DP）进程组中对范数统计量执行归约。 +4. 使用相同的裁剪系数缩放每个进程上的本地梯度分片。 + +启用流水线并行（PP）后，范数统计量还会在 PP 进程组中进一步归约。 + +流水线调度器会在所有微批次执行完成后调用 `ClipGradNormConfigured`。 + +## GPT-2 和 LLaMA-3 配置参数 + +GPT-2 和 LLaMA-3 示例提供了以下命令行选项。 + +### 梯度裁剪阈值 + +```text +--clip_grad_norm=-1 +``` + +负值表示禁用梯度裁剪。 + +### 梯度范数类型 + +```text +--grad_norm_type=2 +``` + +指定梯度范数类型，此处表示使用 2-范数。 + +### 非有限值检查 + +```text +--clip_grad_error_if_nonfinite=true +``` + +当梯度总范数为 `NaN`、`+inf` 或 `-inf` 等非有限值时触发错误。 + +### 多张量缩放模式 + +```text +--clip_grad_foreach=auto|true|false +``` + +控制是否使用多张量批量缩放路径： + +- `auto`：自动选择。 +- `true`：强制启用。 +- `false`：禁用。 + +## 日志记录 + +启用梯度裁剪后，日志中会输出一条 `total_grad_norm` 记录，用于显示裁剪前的梯度总范数。 + +## 并行训练支持 + +对于张量并行（TP）中的复制参数，`DistributedOptimizer` 会根据参数所有权进行过滤，以避免重复统计。 + +完整的 TP/SP/PP/vPP 混合并行训练测试由以下文件中的可选测试矩阵提供： + +```text +scripts/test_config.json +``` diff --git a/gradient_norm_clipping_acceptance_report.md b/gradient_norm_clipping_acceptance_report.md new file mode 100644 index 000000000..0af96990b --- /dev/null +++ b/gradient_norm_clipping_acceptance_report.md @@ -0,0 +1,214 @@ +# 梯度裁剪（Gradient Norm Clipping）最终验收报告 + +## 1. 验收结论 + +梯度裁剪已纳入 `Optimizer` 公共多态接口，完成 CPU/CUDA 范数统计、CPU/CUDA 原地缩放、`DistributedOptimizer` 的 DP/TP/SP/PP/ZeRO 归约、GPT-2/LLaMA3 example 参数和训练调用、测试与文档集成。 + +- CPU/CUDA 裁剪专项：16 个用例，15 passed，1 个 CUDA CPU-only death-test 按设计 skipped，退出码 0。 +- 两进程 CUDA 分布式专项：5/5 passed，退出码 0；覆盖 ZeRO-1/2 shard、empty shard、全局 norm、p=0、p=-∞、`foreach=true`。 +- 单卡 BF16+p=0、2 卡 ZeRO-2+foreach、4 卡 TP4/SP、2 卡 PP2/vPP2+LoRA 真实训练均退出码 0，并输出 `total_grad_norm`。 +- `PROFILE_MODE=ON` 下 `ScaleInplaceMulti=1`、逐 tensor `ScaleInplace=0`，证明 `foreach=true` 使用批量 kernel 路径。 + +## 2. 提交要求对照 + +| 项目 | 结果 | 证据 | +| --- | --- | --- | +| 公共接口与默认语义 | 通过 | `Optimizer::ClipGradNorm_` virtual API、wrapper、配置入口 | +| CPU/CUDA 实现 | 通过 | `ScaleInplace` 与 `ScaleInplaceMulti`，CMake CUDA 构建 | +| 分布式和分片语义 | 通过 | 两进程 5 个专项、2 卡 ZeRO-2 smoke | +| example 集成 | 通过 | GPT-2/LLaMA3 CLI、普通训练和 PP schedule 均调用裁剪 | +| 报告与文档 | 通过 | 本报告、`docs/gradient_norm_clipping.md` | +| 代码提交/PR review/approve | 按要求保留 | 本轮不执行 commit/push/review/approve | + +## 3. 公共 API 与行为 + +```cpp +virtual std::shared_ptr ClipGradNorm_( + const std::vector>& parameters, + float max_norm, + float norm_type = 2.0f, + bool error_if_nonfinite = false, + std::optional foreach = std::nullopt); +``` + +`ClipGradNorm` 返回裁剪前的 CPU FP32 scalar，并原地缩放选中参数的梯度；无梯度参数被忽略，重复参数指针只计算一次，梯度存储不被替换。缩放系数为： + +```text +min(max_norm / (total_norm + 1e-6), 1) +``` + +支持的 `norm_type`： + +- 有限正 p 和有限负 p：按通用 p-norm 公式计算； +- `0`：按 PyTorch 语义，先对每个梯度 tensor 计非零元素，再对 tensor 列表取 0-norm； +- `+∞`：最大绝对值； +- `-∞`：最小绝对值； +- `error_if_nonfinite=true`：在缩放前检查并终止。 + +GPT-2 与 LLaMA3 暴露： + +```text +--clip_grad_norm=-1 +--grad_norm_type=2 +--clip_grad_error_if_nonfinite=true +--clip_grad_foreach=auto|true|false +``` + +## 4. 实现与数据流 + +`infini_train/src/optimizer.cc` 负责重复指针过滤、FP32 统计、p=0/负 p/±∞ 分支、非有限检查和按 device/dtype 分组的缩放。 + +`infini_train/src/kernels/cpu/accumulate_grad.cc`： + +- `ScaleInplace`：单 tensor OpenMP 缩放； +- `ScaleInplaceMulti`：多个 tensor 的批量 OpenMP 调度。 + +`infini_train/src/kernels/cuda/accumulate_grad.cu`： + +- `ScaleInplace`：单 tensor CUDA kernel； +- `ScaleInplaceMulti`：一次 pointer/offset 元数据上传和一次多 tensor CUDA kernel launch，直接写回原始 view，保留 ZeRO flat-buffer alias。 + +`DistributedOptimizer::ClipGradNorm_` 的顺序为： + +1. `FinishGradSync()`，等待梯度 collective； +2. 选择请求参数对应的 shard； +3. TP replicated bias 只在 TP group rank 0 计入 norm； +4. base optimizer 计算本地 shard 统计； +5. DP、TP、PP group 按 norm 类型执行 sum/max/min 归约； +6. 所有 rank 使用同一全局系数批量缩放本地 shard。 + +Pipeline schedule 在所有 micro-batch backward 完成后裁剪一次，并透传最后一次 `total_grad_norm` 到 example 日志。 + +## 5. 测试结果 + +主要测试文件： + +- `tests/optimizer/test_clip_grad_norm.cc` +- `tests/optimizer/test_optimizer_parameter_names.cc` +- `tests/optimizer/CMakeLists.txt` + +覆盖内容： + +- L2、L1、非整数 p、`+∞`； +- p=0、有限负 p、`-∞`； +- `max_norm=0`、空梯度、重复参数、非有限梯度； +- CPU FP16/BF16/FP32 混合 dtype； +- `foreach=true` 批量语义； +- ZeRO local/global/empty shard、TP owner filtering。 + +远程 CUDA 裁剪专项： + +```text +16 tests ran +15 passed +1 skipped (CUDA/ClipGradNormTest.ErrorOnNonFiniteBeforeScaling is CPU-only) +STATUS:0 +``` + +远程两进程分布式专项： + +```text +5 tests ran +5 passed on each rank +STATUS:0 +``` + +通过的专项包括： + +```text +DistributedOptimizerPropagatesNamesToShardOptimizer +DistributedOptimizerClipGradNormUsesZero2LocalShard +DistributedOptimizerClipGradNormHandlesEmptyLocalShard +DistributedOptimizerClipGradNormUsesGlobalShardNorm +DistributedOptimizerClipGradNormSupportsZeroAndNegativeInfinity +``` + +本地 CPU 优化器完整回归： + +```text +36 tests ran, 31 passed, 5 CUDA-only skipped, exit code 0 +``` + +远程 CUDA 优化器完整回归： + +```text +70 tests ran, 59 passed, 11 design-skipped (CPU-only or distributed-only filters), exit code 0 +``` + +## 6. 真实 GPU 训练验收 + +### 单卡 FP32 + foreach + +```text +train loss 5.358113 +total_grad_norm 65.678352 +DP=1, TP=1, SP=1, PP=1 +exit code 0 +``` + +### 单卡 BF16 + p=0 + +```text +train loss 5.342469 +total_grad_norm 149.000000 +DP=1, TP=1, SP=1, PP=1 +exit code 0 +``` + +### 2 卡 ZeRO-2 + foreach=true + +```text +train loss 5.439159 +total_grad_norm 51.047932 +DP=2, TP=1, SP=1, PP=1 +exit code 0 +``` + +### 4 卡 TP4/SP + +```text +train loss 5.356194 +total_grad_norm 18.435957 +DP=1, TP=4, SP=4, PP=1 +exit code 0 +``` + +### 2 卡 PP2/vPP2 + LoRA + +```text +train loss 5.358113 +total_grad_norm 0.896690 +DP=1, TP=1, SP=1, PP=2 +exit code 0 +``` + +PP 日志已补齐 `total_grad_norm`，不再出现 pipeline 路径只有 loss 而没有 norm 的情况。 + +## 7. Foreach/Profile 证据 + +profile 构建命令使用 `-DPROFILE_MODE=ON`，同一轮 GPT-2 FP32 + `clip_grad_foreach=true` 训练结果： + +```text +ScaleInplaceMulti: 1 +ScaleInplace: 0 +GPT2_STATUS:0 +``` + +示例 profile 记录： + +```text +ScaleInplaceMulti Device(CUDA, 0) +``` + +这证明 `foreach=true` 没有退化到逐 tensor `ScaleInplace`。缩放 kernel 仍在 CUDA stream 上提交；范数最终 scalar 的 host copy 属于当前返回 CPU scalar 的 API 行为。 + +## 9. 配置矩阵 + +`scripts/test_config.json` 新增 `gradient_clipping_excellent` 组： + +- `single_fp32_foreach` +- `single_bf16_p0` +- `tp4_sp_zero2` +- `pp2_vpp2_lora` + +每个 case 都配置了 `clip_grad_norm`、`grad_norm_type` 和 `clip_grad_foreach`；GPT-2 和 LLaMA3 测试组入口均已包含该 tag。