diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc index 7697ad213..4af5b8260 100644 --- a/example/gpt2/main.cc +++ b/example/gpt2/main.cc @@ -357,7 +357,6 @@ void Train(const nn::parallel::Rank &rank) { } else { optimizer = optimizer_creator(named_parameters); } - 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); @@ -516,6 +515,7 @@ void Train(const nn::parallel::Rank &rank) { LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish backward"; } + nn::parallel::FinalizeModelGrads({model}); optimizer->Step(); if (scheduler) { scheduler->Step(); diff --git a/example/llama3/main.cc b/example/llama3/main.cc index a9f405fa5..b8342ca90 100644 --- a/example/llama3/main.cc +++ b/example/llama3/main.cc @@ -347,7 +347,6 @@ void Train(const nn::parallel::Rank &rank) { } else { optimizer = optimizer_creator(named_parameters); } - 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); @@ -503,6 +502,7 @@ void Train(const nn::parallel::Rank &rank) { LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish backward"; } + nn::parallel::FinalizeModelGrads({model}); optimizer->Step(); if (scheduler) { scheduler->Step(); diff --git a/example/mixtral/main.cc b/example/mixtral/main.cc index ce4c8cf36..4998675f8 100644 --- a/example/mixtral/main.cc +++ b/example/mixtral/main.cc @@ -22,6 +22,7 @@ #include "infini_train/include/nn/modules/loss.h" #include "infini_train/include/nn/modules/transformer/transformer.h" #include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/nn/parallel/utils.h" #include "infini_train/include/optimizer.h" #include "infini_train/include/tensor.h" #ifdef PROFILE_MODE @@ -150,6 +151,7 @@ int main(int argc, char *argv[]) { auto loss_cpu = loss->To(Device()); lossf += static_cast(loss_cpu.DataPtr())[0]; } + infini_train::nn::parallel::FinalizeModelGrads({model}); optimizer->Step(); device_impl->SynchronizeDevice(train_device); diff --git a/example/qwen3/main.cc b/example/qwen3/main.cc index ba969fd21..11249f9fc 100644 --- a/example/qwen3/main.cc +++ b/example/qwen3/main.cc @@ -502,6 +502,7 @@ void Train(const nn::parallel::Rank &rank) { LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish backward"; } + nn::parallel::FinalizeModelGrads({model}); optimizer->Step(); if (scheduler) { scheduler->Step(); diff --git a/infini_train/include/nn/parallel/ddp/distributed_data_parallel.h b/infini_train/include/nn/parallel/ddp/distributed_data_parallel.h index 905816b56..d5d527610 100644 --- a/infini_train/include/nn/parallel/ddp/distributed_data_parallel.h +++ b/infini_train/include/nn/parallel/ddp/distributed_data_parallel.h @@ -34,6 +34,8 @@ class DistributedDataParallel : public nn::Module { std::unique_ptr no_sync() override; + void FinishGradSync(); + DistributedDataParallelConfig ddp_config() const { return ddp_config_; } const std::vector> ¶m_grad_buffers() const { return param_grad_buffers_; } diff --git a/infini_train/include/nn/parallel/ddp/distributed_optimizer.h b/infini_train/include/nn/parallel/ddp/distributed_optimizer.h index d7cea198e..b27e6f7b0 100644 --- a/infini_train/include/nn/parallel/ddp/distributed_optimizer.h +++ b/infini_train/include/nn/parallel/ddp/distributed_optimizer.h @@ -37,9 +37,6 @@ class DistributedOptimizer final : public infini_train::Optimizer { void LoadStateDict(const std::unordered_map> &state_dict) override; - void StartGradSync(); - void FinishGradSync(); - void StartParamSync(bool force_sync = false); void FinishParamSync(bool skip_next_bucket_dispatch = false); 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..a7866c4b3 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 @@ -5,6 +5,7 @@ #include #include #include +#include #include #include "infini_train/include/datatype.h" @@ -19,6 +20,9 @@ class Work; } // namespace infini_train namespace infini_train::nn::parallel { +// Original parameter and its local gradient view, registered by DistributedOptimizer. +using LocalGradShard = std::pair, std::shared_ptr>; + class ParamAndGradBucket { public: /** @@ -123,6 +127,10 @@ class ParamAndGradBucketGroup { // ZeRO-2: Get a bucket's local grad shard buffer std::shared_ptr GetLocalGradShardBuffer(size_t bucket_idx) const; + const std::vector &local_grad_shards() const; + + void set_local_grad_shards(std::vector shards); + const DistributedDataParallelConfig &config() const { return ddp_config_; } private: @@ -146,6 +154,8 @@ class ParamAndGradBucketGroup { // ZeRO-2: persistent grad shard buffers and temporary full grad buffers std::vector> grad_shard_buffer_list_; std::vector> temp_full_grad_buffer_list_; + // These views share optimizer grad storage and persist across iteration resets. + std::vector local_grad_shards_; std::shared_ptr next_param_gather_bucket_group_ = nullptr; diff --git a/infini_train/include/nn/parallel/utils.h b/infini_train/include/nn/parallel/utils.h index 8aa11856a..1d3fe975c 100644 --- a/infini_train/include/nn/parallel/utils.h +++ b/infini_train/include/nn/parallel/utils.h @@ -7,6 +7,9 @@ namespace infini_train { class Tensor; +namespace nn { +class Module; +} } // namespace infini_train namespace infini_train::nn::parallel { @@ -31,4 +34,6 @@ std::vector> GatherFromSPRegionFunc(const std::shared_pt std::vector> ScatterToTPRegionFunc(const std::shared_ptr &input); std::vector> ReduceFromTPRegionFunc(const std::shared_ptr &input); std::vector> CopyToTPRegionFunc(const std::shared_ptr &input); + +void FinalizeModelGrads(const std::vector> &model_chunks); } // namespace infini_train::nn::parallel diff --git a/infini_train/include/tensor.h b/infini_train/include/tensor.h index dcfd8927f..3aabe28f1 100644 --- a/infini_train/include/tensor.h +++ b/infini_train/include/tensor.h @@ -80,6 +80,9 @@ class Tensor : public std::enable_shared_from_this { size_t NumElements() const; DataType Dtype() const; + void set_sequence_parallel(bool enabled); + bool sequence_parallel() const; + std::shared_ptr Detach() const; void Fill(Scalar value); @@ -242,6 +245,7 @@ class Tensor : public std::enable_shared_from_this { private: std::shared_ptr grad_ = nullptr; bool requires_grad_ = false; + bool sequence_parallel_ = false; bool is_leaf_ = true; std::shared_ptr grad_fn_ = nullptr; int output_idx_ = 0; diff --git a/infini_train/src/nn/modules/normalization.cc b/infini_train/src/nn/modules/normalization.cc index 388b04de5..4bd0b0dab 100644 --- a/infini_train/src/nn/modules/normalization.cc +++ b/infini_train/src/nn/modules/normalization.cc @@ -7,6 +7,7 @@ #include "infini_train/include/device.h" #include "infini_train/include/nn/functional.h" #include "infini_train/include/nn/init.h" +#include "infini_train/include/nn/parallel/global.h" #include "infini_train/include/tensor.h" namespace infini_train::nn { @@ -18,6 +19,10 @@ LayerNorm::LayerNorm(const std::vector &normalized_shape, float eps, De = std::make_shared(normalized_shape, DataType::kFLOAT32, device_)->RequiresGrad(); parameters_[kParamBiasName] = std::make_shared(normalized_shape, DataType::kFLOAT32, device_)->RequiresGrad(); + if (parallel::global::GetSequenceParallelEnabled()) { + parameters_[kParamWeightName]->set_sequence_parallel(true); + parameters_[kParamBiasName]->set_sequence_parallel(true); + } ResetParameters(); } @@ -35,6 +40,9 @@ void LayerNorm::ResetParameters() { RMSNorm::RMSNorm(int64_t dim, float eps, Device device) : CloneableModule(kType), eps_(eps) { parameters_[kParamWeightName] = std::make_shared(std::vector{dim}, DataType::kFLOAT32, device)->RequiresGrad(); + if (parallel::global::GetSequenceParallelEnabled()) { + parameters_[kParamWeightName]->set_sequence_parallel(true); + } nn::init::Ones(parameters_[kParamWeightName]); } diff --git a/infini_train/src/nn/modules/transformer/causal_self_attention.cc b/infini_train/src/nn/modules/transformer/causal_self_attention.cc index 9c32fa23d..31676b1b4 100644 --- a/infini_train/src/nn/modules/transformer/causal_self_attention.cc +++ b/infini_train/src/nn/modules/transformer/causal_self_attention.cc @@ -26,6 +26,10 @@ CausalSelfAttention::CausalSelfAttention(const TransformerConfig &config) : Clon if (config_.qk_layernorm) { modules_[kQNormLayerName] = std::make_shared(head_dim_, config_.norm_eps); modules_[kKNormLayerName] = std::make_shared(head_dim_, config_.norm_eps); + // NOTE(zbl): In Qwen3-8B, Q/K norm sees full sequences and TP-local heads. + // So we only need to finalize its replicated weights across TP regardless of SP. + modules_[kQNormLayerName]->parameter(RMSNorm::kParamWeightName)->set_sequence_parallel(false); + modules_[kKNormLayerName]->parameter(RMSNorm::kParamWeightName)->set_sequence_parallel(false); } int64_t qkv_dim = (config.n_head + 2 * n_kv_head_) * head_dim_; diff --git a/infini_train/src/nn/modules/transformer/moe/router.cc b/infini_train/src/nn/modules/transformer/moe/router.cc index 252086846..b045400aa 100644 --- a/infini_train/src/nn/modules/transformer/moe/router.cc +++ b/infini_train/src/nn/modules/transformer/moe/router.cc @@ -11,6 +11,7 @@ #include "infini_train/include/nn/functional.h" #include "infini_train/include/nn/init.h" #include "infini_train/include/nn/modules/transformer/moe/moe_utils.h" +#include "infini_train/include/nn/parallel/global.h" #include "infini_train/include/tensor.h" namespace infini_train::nn::moe { @@ -24,12 +25,18 @@ TopKRouter::TopKRouter(const TransformerConfig &config) : CloneableModule(kType) = std::make_shared(std::vector{moe_config.num_experts, config_.n_embd}, DataType::kFLOAT32, device_) ->RequiresGrad(); + if (parallel::global::GetSequenceParallelEnabled()) { + parameters_[kParamWeightName]->set_sequence_parallel(true); + } init::KaimingUniform(parameters_[kParamWeightName]); if (config_.add_bias_linear) { parameters_[kParamBiasName] = std::make_shared(std::vector{moe_config.num_experts}, DataType::kFLOAT32, device_) ->RequiresGrad(); + if (parallel::global::GetSequenceParallelEnabled()) { + parameters_[kParamBiasName]->set_sequence_parallel(true); + } parameters_[kParamBiasName]->Fill(0.0f); } } diff --git a/infini_train/src/nn/modules/transformer/transformer.cc b/infini_train/src/nn/modules/transformer/transformer.cc index 99a739d2d..f1048058a 100644 --- a/infini_train/src/nn/modules/transformer/transformer.cc +++ b/infini_train/src/nn/modules/transformer/transformer.cc @@ -32,6 +32,9 @@ TransformerFirstStage::TransformerFirstStage(const TransformerConfig &config) // Only learned absolute position embedding uses a trainable WPE table. if (config_.position_embedding_type == PositionEmbeddingType::kLearnedAbsolute) { modules_[kWPELayerName] = std::make_shared(config_.block_size, config_.n_embd); + if (parallel::global::GetSequenceParallelEnabled()) { + modules_[kWPELayerName]->parameter(Embedding::kParamWeightName)->set_sequence_parallel(true); + } } else if (config_.position_embedding_type != PositionEmbeddingType::kRoPE) { LOG(FATAL) << "Unsupported position embedding type"; } diff --git a/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc b/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc index dd05a8d71..e063fcb6f 100644 --- a/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc +++ b/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc @@ -231,6 +231,10 @@ std::unique_ptr DistributedDataParallel::no_sync() { return std::make_unique([this, previous] { SetIsLastMicrobatch(previous); }); } +void DistributedDataParallel::FinishGradSync() { + for (auto &group : bucket_groups_) { group->FinishGradSync(); } +} + void DistributedDataParallel::SetIsLastMicrobatch(bool is_last_microbatch) { is_last_microbatch_->store(is_last_microbatch, std::memory_order_relaxed); if (reducer_) { diff --git a/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc b/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc index 523bcf2d7..a8fa73e64 100644 --- a/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc +++ b/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc @@ -68,6 +68,7 @@ void DistributedOptimizer::BuildShardParamsAndBindGrads(const AddShardParam &add size_t num_shard_params = 0; for (const auto &group : bucket_groups_) { + std::vector local_grad_shards; const bool use_grad_shard = group->config().zero_stage >= 2; const auto &buckets = group->buckets(); for (size_t bucket_idx = 0; bucket_idx < buckets.size(); ++bucket_idx) { @@ -107,6 +108,7 @@ void DistributedOptimizer::BuildShardParamsAndBindGrads(const AddShardParam &add auto param_piece = std::make_shared(*bucket_param, param_piece_offset_bytes, std::vector{static_cast(piece_numel)}); + param_piece->set_sequence_parallel(param->sequence_parallel()); auto grad_piece = std::make_shared(*bucket_grad, grad_piece_offset_bytes, std::vector{static_cast(piece_numel)}); @@ -115,24 +117,18 @@ void DistributedOptimizer::BuildShardParamsAndBindGrads(const AddShardParam &add // NOTE(zbl): Do not call `param->set_grad(grad_piece);` under ZeRO-2. // The base optimizer updates param_piece views only; original param->grad() // would be a partial flattened shard and does not represent the full parameter grad. + local_grad_shards.emplace_back(param, grad_piece); add_shard_param(param, param_piece); ++num_shard_params; } } + group->set_local_grad_shards(std::move(local_grad_shards)); } CHECK_GT(num_shard_params, 0) << "DistributedOptimizer: this DP rank owns no param pieces. " << "Check bucket padding/divisibility and param bucketing order."; } -void DistributedOptimizer::StartGradSync() { - for (auto &group : bucket_groups_) { group->StartGradSync(); } -} - -void DistributedOptimizer::FinishGradSync() { - for (auto &group : bucket_groups_) { group->FinishGradSync(); } -} - void DistributedOptimizer::StartParamSync(bool force_sync) { for (auto &group : bucket_groups_) { group->StartParamSync(force_sync); } } @@ -171,14 +167,10 @@ float DistributedOptimizer::learning_rate() const { } void DistributedOptimizer::Step() { - // 1. Ensure grads are synced - FinishGradSync(); - - // 2. Base optimizer step on owned param pieces CHECK(base_optimizer_) << "DistributedOptimizer: base optimizer is null."; base_optimizer_->Step(); - // 3. Gather updated param shards back to full params + // Gather updated param shards back to full params StartParamSync(/*force_sync=*/false); // TODO(zbl): Delay sync call until param is actually used in next step FinishParamSync(/*skip_next_bucket_dispatch=*/true); diff --git a/infini_train/src/nn/parallel/ddp/param_and_grad_buffer.cc b/infini_train/src/nn/parallel/ddp/param_and_grad_buffer.cc index 8a3c2d052..970c30f27 100644 --- a/infini_train/src/nn/parallel/ddp/param_and_grad_buffer.cc +++ b/infini_train/src/nn/parallel/ddp/param_and_grad_buffer.cc @@ -240,6 +240,12 @@ std::shared_ptr ParamAndGradBucketGroup::GetLocalGradShardBuffer(size_t return grad_shard_buffer_list_[bucket_idx]; } +const std::vector &ParamAndGradBucketGroup::local_grad_shards() const { return local_grad_shards_; } + +void ParamAndGradBucketGroup::set_local_grad_shards(std::vector shards) { + local_grad_shards_ = std::move(shards); +} + void ParamAndGradBucketGroup::StartGradSync() { if (!collective_pg_) { LOG(FATAL) << "ParamAndGradBucketGroup: StartGradSync() called with null collective_pg_."; diff --git a/infini_train/src/nn/parallel/pp/pipeline_schedule.cc b/infini_train/src/nn/parallel/pp/pipeline_schedule.cc index 0ea2a4710..40ba16ddf 100644 --- a/infini_train/src/nn/parallel/pp/pipeline_schedule.cc +++ b/infini_train/src/nn/parallel/pp/pipeline_schedule.cc @@ -15,6 +15,7 @@ #include "infini_train/include/nn/parallel/global.h" #include "infini_train/include/nn/parallel/pp/pipeline_stage.h" #include "infini_train/include/nn/parallel/pp/send_recv.h" +#include "infini_train/include/nn/parallel/utils.h" #include "infini_train/include/optimizer.h" #include "infini_train/include/tensor.h" @@ -299,6 +300,7 @@ float PipelineSchedule::Step(std::shared_ptr input, std::shared_ptrchunks()); optimizer->Step(); return lossf; diff --git a/infini_train/src/nn/parallel/tensor_parallel.cc b/infini_train/src/nn/parallel/tensor_parallel.cc index b16c526e6..65c55bbba 100644 --- a/infini_train/src/nn/parallel/tensor_parallel.cc +++ b/infini_train/src/nn/parallel/tensor_parallel.cc @@ -305,6 +305,9 @@ RowParallelLinear::RowParallelLinear(int64_t in_features, int64_t out_features, if (bias) { parameters_[kParamBiasName] = std::make_shared(std::vector{out_features}, DataType::kFLOAT32, device_)->RequiresGrad(); + if (sequence_parallel_) { + parameters_[kParamBiasName]->set_sequence_parallel(true); + } } LinearResetParameters(parameters_[kParamWeightName], bias ? parameters_[kParamBiasName] : nullptr); diff --git a/infini_train/src/nn/parallel/utils.cc b/infini_train/src/nn/parallel/utils.cc index ee28a1694..9d3e871ec 100644 --- a/infini_train/src/nn/parallel/utils.cc +++ b/infini_train/src/nn/parallel/utils.cc @@ -1,13 +1,66 @@ #include "infini_train/include/nn/parallel/utils.h" +#include + #include "glog/logging.h" #include "infini_train/include/nn/functional.h" +#include "infini_train/include/nn/modules/module.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/process_group.h" +#include "infini_train/include/nn/parallel/reduce_op_type.h" #include "infini_train/include/tensor.h" namespace infini_train::nn::parallel { +namespace { +const ProcessGroup *GetTensorParallelGroup(const Tensor &tensor) { + const int global_rank = tensor.GetDevice().Rank().GlobalRank(); + return ProcessGroupFactory::Instance(tensor.GetDevice().type()) + ->Get(GetTensorParallelProcessGroupName(global_rank)); +} + +void FinalizeSequenceParallelGradients(const std::vector ¶m_grads) { + // SP replicas see different sequence shards, so replicated parameter grads + // must be summed across TP before the optimizer consumes them. + if (!global::GetSequenceParallelEnabled() || global::GetTensorParallelSize() <= 1 || param_grads.empty()) { + return; + } + + const ProcessGroup *tp_group = nullptr; + for (const auto &[param, grad] : param_grads) { + if (!param || !param->sequence_parallel() || !grad) { + continue; + } + + if (tp_group == nullptr) { + tp_group = GetTensorParallelGroup(*param); + CHECK_NOTNULL(tp_group); + } + tp_group->AllReduce(grad, function::ReduceOpType::kSum, /*async_op=*/false); + } +} + +void FinalizeQKNormWeightGradients(const nn::Module &model_chunk, const std::vector ¶m_grads) { + // CausalSelfAttention Q/K norm weights (e.g. in Qwen3) are shared across heads and replicated across TP ranks. + // Each rank computes gradients for TP-local heads, so sum them across TP regardless of SP. + // These weights have sequence_parallel=false and are not reduced by FinalizeSequenceParallelGradients. + std::unordered_set qk_norm_params; + for (const auto &[name, param] : model_chunk.NamedParameters()) { + if (name.ends_with("q_norm.weight") || name.ends_with("k_norm.weight")) { + qk_norm_params.insert(param.get()); + } + } + + for (const auto &[param, grad] : param_grads) { + if (param && param->requires_grad() && grad && qk_norm_params.contains(param.get())) { + const auto *tp_group = GetTensorParallelGroup(*param); + CHECK_NOTNULL(tp_group); + tp_group->AllReduce(grad, function::ReduceOpType::kSum, /*async_op=*/false); + } + } +} +} // namespace std::string GetDataParallelProcessGroupName(int global_rank) { return "DP" + std::to_string(global::GetGroupId(global::DP, global_rank)); @@ -60,4 +113,42 @@ std::shared_ptr GatherTensorParallelShard(const std::shared_ptr return nn::function::Concat(rank_major_shards, dim)->Contiguous(); } +// Call once after all microbatch backwards, before gradient clipping and optimizer updates. +void FinalizeModelGrads(const std::vector> &model_chunks) { + // Finish DP gradient synchronization for all chunks before TP reduction. + for (const auto &model_chunk : model_chunks) { + if (auto ddp = std::dynamic_pointer_cast(model_chunk)) { + ddp->FinishGradSync(); + } + } + + if (global::GetTensorParallelSize() <= 1) { + // DP synchronization is complete; skip TP-specific finalization. + return; + } + + for (const auto &model_chunk : model_chunks) { + // Pair original parameters with the gradient views consumed by the optimizer. + std::vector param_grads; + auto ddp = std::dynamic_pointer_cast(model_chunk); + if (ddp && ddp->ddp_config().zero_stage >= 1) { + // ZeRO-1/2 use DP-local gradient shard views registered by DistOpt. + for (const auto &group : ddp->bucket_groups()) { + const auto &shards = group->local_grad_shards(); + param_grads.insert(param_grads.end(), shards.begin(), shards.end()); + } + } else { + // Ordinary optimizer / ZeRO-0 use full gradients from param.grad(). + for (const auto ¶m : model_chunk->Parameters()) { param_grads.emplace_back(param, param->grad()); } + } + // TP SUM combines local-token gradients for SP norms, row bias, position embeddings and routers. + FinalizeSequenceParallelGradients(param_grads); + + // TP SUM combines local-head gradients for replicated Q/K norm weights regardless of SP. + FinalizeQKNormWeightGradients(*model_chunk, param_grads); + } + + // NOTE(zbl): Extend this entry for PP tied embeddings, MoE shared parameters and loss normalization. +} + } // namespace infini_train::nn::parallel diff --git a/infini_train/src/tensor.cc b/infini_train/src/tensor.cc index 4e61e221b..70d9591af 100644 --- a/infini_train/src/tensor.cc +++ b/infini_train/src/tensor.cc @@ -57,6 +57,7 @@ Tensor::Tensor(const Tensor &tensor, size_t offset, const std::vector & : buffer_(tensor.buffer_), offset_(tensor.offset_ + offset), dims_(dims), num_elements_(std::accumulate(dims.begin(), dims.end(), 1, std::multiplies())), dtype_(tensor.dtype_) { CHECK_LE(offset_ + kDataTypeToSize.at(dtype_) * num_elements_, buffer_->Size()); + sequence_parallel_ = tensor.sequence_parallel_; } Tensor::Tensor(const float *data, const std::vector &dims, DataType dtype, Device device) @@ -104,6 +105,10 @@ size_t Tensor::NumElements() const { return num_elements_; } DataType Tensor::Dtype() const { return dtype_; } +void Tensor::set_sequence_parallel(bool enabled) { sequence_parallel_ = enabled; } + +bool Tensor::sequence_parallel() const { return sequence_parallel_; } + std::shared_ptr Tensor::Detach() const { return std::make_shared(*this, 0, dims_); } void Tensor::Fill(Scalar value) { @@ -128,9 +133,11 @@ Eigen::Map> Tensor::Eig Tensor Tensor::To(Device device) { const auto buffer_device = buffer_->GetDevice(); if (device == buffer_device) { - auto new_tensor = Tensor(*this, offset_, dims_); + auto new_tensor = Tensor(*this, 0, dims_); + new_tensor.requires_grad_ = requires_grad_; + new_tensor.sequence_parallel_ = sequence_parallel_; if (grad_) { - new_tensor.grad_ = std::make_unique(*grad_.get(), grad_->offset_, grad_->dims_); + new_tensor.grad_ = std::make_unique(*grad_.get(), 0, grad_->dims_); } return new_tensor; } @@ -170,15 +177,18 @@ Tensor Tensor::To(Device device) { } new_tensor.requires_grad_ = requires_grad_; + new_tensor.sequence_parallel_ = sequence_parallel_; return new_tensor; } Tensor Tensor::To(DataType dtype) { if (dtype == dtype_) { - auto new_tensor = Tensor(*this, offset_, dims_); + auto new_tensor = Tensor(*this, 0, dims_); + new_tensor.requires_grad_ = requires_grad_; + new_tensor.sequence_parallel_ = sequence_parallel_; if (grad_) { - new_tensor.grad_ = std::make_unique(*grad_.get(), grad_->offset_, grad_->dims_); + new_tensor.grad_ = std::make_unique(*grad_.get(), 0, grad_->dims_); } return new_tensor; } @@ -194,6 +204,7 @@ Tensor Tensor::To(DataType dtype) { } new_tensor.requires_grad_ = requires_grad_; + new_tensor.sequence_parallel_ = sequence_parallel_; return new_tensor; } diff --git a/tests/tensor/test_tensor_copy.cc b/tests/tensor/test_tensor_copy.cc index 88d597b8e..84dcbc434 100644 --- a/tests/tensor/test_tensor_copy.cc +++ b/tests/tensor/test_tensor_copy.cc @@ -27,4 +27,34 @@ TEST_P(TensorCopyTest, CopiesPreservesDataType) { EXPECT_EQ(target->Dtype(), DataType::kFLOAT32); } +TEST_P(TensorCopyTest, NoOpConversionsPreserveParameterMetadataAndViews) { + auto storage = std::make_shared(std::vector{12}, DataType::kFLOAT32, GetDevice()); + auto grad_storage = std::make_shared(std::vector{12}, DataType::kFLOAT32, GetDevice()); + auto parameter = std::make_shared(*storage, 4 * sizeof(float), std::vector{4}); + auto grad = std::make_shared(*grad_storage, 4 * sizeof(float), std::vector{4}); + parameter->RequiresGrad(); + parameter->set_sequence_parallel(true); + parameter->set_grad(grad); + + for (auto converted : {parameter->To(GetDevice()), parameter->To(parameter->Dtype())}) { + EXPECT_TRUE(converted.requires_grad()); + EXPECT_TRUE(converted.sequence_parallel()); + EXPECT_EQ(converted.DataPtr(), parameter->DataPtr()); + EXPECT_EQ(converted.Dims(), parameter->Dims()); + ASSERT_NE(converted.grad(), nullptr); + EXPECT_EQ(converted.grad()->DataPtr(), grad->DataPtr()); + EXPECT_EQ(converted.grad()->Dims(), grad->Dims()); + } +} + +TEST_P(TensorCopyTest, NoOpConversionsPreserveFrozenTensorMetadata) { + auto tensor = std::make_shared(std::vector{4}, DataType::kFLOAT32, GetDevice()); + for (auto converted : {tensor->To(GetDevice()), tensor->To(tensor->Dtype())}) { + EXPECT_FALSE(converted.requires_grad()); + EXPECT_FALSE(converted.sequence_parallel()); + EXPECT_EQ(converted.grad(), nullptr); + EXPECT_EQ(converted.DataPtr(), tensor->DataPtr()); + } +} + INFINI_TRAIN_REGISTER_TEST(TensorCopyTest);