diff --git a/example/llama3/checkpoint_loader.cc b/example/llama3/checkpoint_loader.cc index f3590af6e..32b11e430 100644 --- a/example/llama3/checkpoint_loader.cc +++ b/example/llama3/checkpoint_loader.cc @@ -159,7 +159,7 @@ std::shared_ptr LoadFromLLMC(const std::string &filepath) // RowParallel (proj) const int64_t in_pp = static_cast(n_embd) / tp_size; - // nn::MLP: c_fc/c_fc2（shard along row），c_proj（shard along col） + // nn::MLP: packed c_fc [gate | up] (shard each block along row), c_proj (shard along col) const int64_t fc_out = ffn_hidden; const int64_t fc_pp = fc_out / tp_size; const int64_t in_fc_pp = ffn_hidden / tp_size; @@ -269,7 +269,7 @@ std::shared_ptr LoadFromLLMC(const std::string &filepath) } } - // transformer.h.{i}.mlp.c_fc.weight : ColumnParallelLinear, but actually applies on "rows" + // transformer.h.{i}.mlp.c_fc.weight (up) -> local packed c_fc rows [fc_pp : 2*fc_pp) local_layer_index = 0; for (int i = 0; i < static_cast(n_layer); ++i) { if (owned_layers[i]) { @@ -277,7 +277,8 @@ std::shared_ptr LoadFromLLMC(const std::string &filepath) nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index), nn::TransformerLayer::kMlpLayerName, nn::MLP::kCFcLayerName, nn::parallel::ColumnParallelLinear::kParamWeightName)]; - ReadMatrixRowShardFloat(ifs, static_cast(tensor->DataPtr()), + float *dst = static_cast(tensor->DataPtr()) + fc_pp * n_embd; + ReadMatrixRowShardFloat(ifs, dst, /*rows=*/fc_out, /*cols=*/n_embd, /*row_start=*/tp_rank * fc_pp, /*row_cnt=*/fc_pp); ++local_layer_index; @@ -287,13 +288,13 @@ std::shared_ptr LoadFromLLMC(const std::string &filepath) } } - // transformer.h.{i}.mlp.c_fc2.weight : ColumnParallelLinear, but actually applies on "rows" + // transformer.h.{i}.mlp.c_fc2.weight (gate) -> local packed c_fc rows [0 : fc_pp) local_layer_index = 0; for (int i = 0; i < static_cast(n_layer); ++i) { if (owned_layers[i]) { auto &tensor = state_dict[std::format("{}.{}.{}.{}.{}.{}", nn::TransformerModel::kTransformerModelName, nn::TransformerChunk::kHLayerName, std::to_string(local_layer_index), - nn::TransformerLayer::kMlpLayerName, nn::MLP::kCFc2LayerName, + nn::TransformerLayer::kMlpLayerName, nn::MLP::kCFcLayerName, nn::parallel::ColumnParallelLinear::kParamWeightName)]; ReadMatrixRowShardFloat(ifs, static_cast(tensor->DataPtr()), /*rows=*/fc_out, /*cols=*/n_embd, diff --git a/example/mixtral/checkpoint_loader.cc b/example/mixtral/checkpoint_loader.cc index c6c8471be..c451ce3cd 100644 --- a/example/mixtral/checkpoint_loader.cc +++ b/example/mixtral/checkpoint_loader.cc @@ -119,10 +119,10 @@ std::shared_ptr LoadFromLLMC(const std::string &filepath, CHECK(ifs) << "Failed to read tensor " << name; }; - auto read_projection_into_packed_qkv = [&](const std::string &packed_qkv_name, int64_t row_offset, int64_t num_rows, - const std::string &projection_name) { - CHECK(state.contains(packed_qkv_name)) << "Model state_dict does not contain " << packed_qkv_name; - std::shared_ptr tensor = state.at(packed_qkv_name); + auto read_projection_into_packed_weight = [&](const std::string &packed_weight_name, int64_t row_offset, + int64_t num_rows, const std::string &projection_name) { + CHECK(state.contains(packed_weight_name)) << "Model state_dict does not contain " << packed_weight_name; + std::shared_ptr tensor = state.at(packed_weight_name); CHECK(tensor->Dtype() == infini_train::DataType::kFLOAT32) << "Only float32 tiny Mixtral LLMC files are supported: " << projection_name; CHECK_EQ(tensor->Dims().size(), 2); @@ -144,17 +144,21 @@ std::shared_ptr LoadFromLLMC(const std::string &filepath, const int64_t head_dim = config.n_embd / config.n_head; const int64_t q_rows = config.n_head * head_dim; const int64_t kv_rows = config.n_kv_head * head_dim; - read_projection_into_packed_qkv(c_attn_name, 0, q_rows, c_attn_name + ".q_proj"); - read_projection_into_packed_qkv(c_attn_name, q_rows, kv_rows, c_attn_name + ".k_proj"); - read_projection_into_packed_qkv(c_attn_name, q_rows + kv_rows, kv_rows, c_attn_name + ".v_proj"); + read_projection_into_packed_weight(c_attn_name, 0, q_rows, c_attn_name + ".q_proj"); + read_projection_into_packed_weight(c_attn_name, q_rows, kv_rows, c_attn_name + ".k_proj"); + read_projection_into_packed_weight(c_attn_name, q_rows + kv_rows, kv_rows, c_attn_name + ".v_proj"); read_tensor_by_state_key(prefix + ".attn.c_proj.weight"); read_tensor_by_state_key(prefix + ".ln_2.weight"); read_tensor_by_state_key(prefix + ".mlp.router.weight"); for (int64_t expert = 0; expert < moe_config.num_experts; ++expert) { const std::string expert_prefix = prefix + ".mlp.experts.expert_" + std::to_string(expert); - read_tensor_by_state_key(expert_prefix + ".c_fc2.weight"); // Mixtral w1/gate_proj - read_tensor_by_state_key(expert_prefix + ".c_fc.weight"); // Mixtral w3/up_proj - read_tensor_by_state_key(expert_prefix + ".c_proj.weight"); // Mixtral w2/down_proj + const std::string packed_fc1_name = expert_prefix + ".c_fc.weight"; + read_projection_into_packed_weight(packed_fc1_name, 0, moe_config.moe_ffn_hidden_size, + expert_prefix + ".c_fc2.weight"); // Mixtral w1/gate_proj + read_projection_into_packed_weight(packed_fc1_name, moe_config.moe_ffn_hidden_size, + moe_config.moe_ffn_hidden_size, + expert_prefix + ".c_fc.weight"); // Mixtral w3/up_proj + read_tensor_by_state_key(expert_prefix + ".c_proj.weight"); // Mixtral w2/down_proj } } read_tensor_by_state_key("transformer.ln_f.weight"); diff --git a/infini_train/include/autograd/activations.h b/infini_train/include/autograd/activations.h index a63977263..079db0fbe 100644 --- a/infini_train/include/autograd/activations.h +++ b/infini_train/include/autograd/activations.h @@ -21,4 +21,16 @@ class Sigmoid : public Function { const std::vector> &output_tensors) override; std::vector> Backward(const std::vector> &grad_outputs) override; }; + +class SwiGLU : public Function { +public: + static constexpr char kType[] = "SwiGLUFunction"; + + SwiGLU() : Function(kType) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; +}; } // namespace infini_train::autograd diff --git a/infini_train/include/common/cuda/common_cuda.h b/infini_train/include/common/cuda/common_cuda.h index 862f41820..b349c9d9c 100644 --- a/infini_train/include/common/cuda/common_cuda.h +++ b/infini_train/include/common/cuda/common_cuda.h @@ -11,6 +11,19 @@ namespace infini_train::common::cuda { +inline size_t ChooseBlockSize(size_t num_elements) { + if (num_elements < 1024) { + return 64; + } + if (num_elements < 65536) { + return 128; + } + if (num_elements < 1048576) { + return 256; + } + return 512; +} + // Common CUDA Macros #define CUDA_CHECK(call) \ do { \ diff --git a/infini_train/include/nn/lora/lora_utils.h b/infini_train/include/nn/lora/lora_utils.h index 89fabbb84..2872bb96d 100644 --- a/infini_train/include/nn/lora/lora_utils.h +++ b/infini_train/include/nn/lora/lora_utils.h @@ -28,6 +28,14 @@ std::shared_ptr SlicePackedQKVRowsForTensorParallel(const std::shared_pt std::shared_ptr RestorePackedQKVRowsFromTensorParallel(const std::shared_ptr &gathered_tensor, int64_t q_rows, int tp_size); +// Internal helper for packed SwiGLU tensors stored as [gate | up] along dim 0. +std::shared_ptr SlicePackedSwiGLURowsForTensorParallel(const std::shared_ptr &full_tensor, int tp_rank, + int tp_size); + +// Internal helper for TP-gathered packed SwiGLU shards stored rank-major as [gate_i | up_i]. +std::shared_ptr RestorePackedSwiGLURowsFromTensorParallel(const std::shared_ptr &gathered_tensor, + int tp_size); + } // namespace detail /** @@ -121,8 +129,8 @@ void LoadLoRAStateDict(std::shared_ptr model, void SaveLoRAWeights(const std::shared_ptr &model, const std::string &filepath); /** - * Load LoRA parameters from file. Packed QKV LoRA-B tensors are split as - * [Qi | Ki | Vi] for the current TP rank. + * Load LoRA parameters from file. Packed QKV and SwiGLU LoRA-B tensors are + * reordered into their rank-local packed layouts. */ void LoadLoRAWeights(std::shared_ptr model, const std::string &filepath); diff --git a/infini_train/include/nn/modules/activations.h b/infini_train/include/nn/modules/activations.h index deb029576..549d93755 100644 --- a/infini_train/include/nn/modules/activations.h +++ b/infini_train/include/nn/modules/activations.h @@ -30,6 +30,7 @@ class SwiGLU : public CloneableModule { static constexpr char kType[] = "SwiGLU"; SwiGLU() : CloneableModule(kType) {} + // The last input dimension is packed as [gate, up], matching Megatron-LM. std::vector> Forward(const std::vector> &x) override; }; } // namespace infini_train::nn diff --git a/infini_train/include/nn/modules/transformer/mlp.h b/infini_train/include/nn/modules/transformer/mlp.h index bb096b7c3..ecf5672b6 100644 --- a/infini_train/include/nn/modules/transformer/mlp.h +++ b/infini_train/include/nn/modules/transformer/mlp.h @@ -13,8 +13,7 @@ class MLP : public infini_train::nn::CloneableModule { static constexpr char kGeluLayerName[] = "gelu"; static constexpr char kCProjLayerName[] = "c_proj"; - static constexpr char kCFc2LayerName[] = "c_fc2"; - static constexpr char kSiluLayerName[] = "silu"; + static constexpr char kSwiGLULayerName[] = "swiglu"; explicit MLP(const TransformerConfig &config); diff --git a/infini_train/src/autograd/activations.cc b/infini_train/src/autograd/activations.cc index bb8b8e5ea..6894788a8 100644 --- a/infini_train/src/autograd/activations.cc +++ b/infini_train/src/autograd/activations.cc @@ -30,4 +30,30 @@ std::vector> Sigmoid::Backward(const std::vectorGetDevice().type(); return {Dispatcher::Instance().Call>({device, "SigmoidBackward"}, output, grad_output)}; } + +std::vector> SwiGLU::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 1); + const auto &input = input_tensors[0]; + CHECK_GT(input->Dims().size(), 0); + CHECK_EQ(input->Dims().back() % 2, 0) << "SwiGLU expects an even last dimension"; + + auto device = input->GetDevice().type(); + return {Dispatcher::Instance().Call>({device, "SwiGLUForward"}, input)}; +} + +void SwiGLU::SetupContext(const std::vector> &input_tensors, + const std::vector> &) { + ctx_.SaveForBackward({input_tensors[0]}); +} + +std::vector> SwiGLU::Backward(const std::vector> &grad_outputs) { + auto saved_tensors = ctx_.GetSavedTensors(); + CHECK_EQ(saved_tensors.size(), 1); + CHECK_EQ(grad_outputs.size(), 1); + const auto &input = saved_tensors[0]; + const auto &grad_output = grad_outputs[0]; + + auto device = input->GetDevice().type(); + return {Dispatcher::Instance().Call>({device, "SwiGLUBackward"}, input, grad_output)}; +} } // namespace infini_train::autograd diff --git a/infini_train/src/kernels/cpu/swiglu.cc b/infini_train/src/kernels/cpu/swiglu.cc new file mode 100644 index 000000000..c87500c06 --- /dev/null +++ b/infini_train/src/kernels/cpu/swiglu.cc @@ -0,0 +1,74 @@ +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +std::shared_ptr SwiGLUForward(const std::shared_ptr &input) { + CHECK(input->Dtype() == DataType::kFLOAT32); + CHECK(input->IsContiguous()); + auto output_dims = input->Dims(); + CHECK_GT(output_dims.size(), 0); + CHECK_EQ(output_dims.back() % 2, 0); + const int64_t hidden = output_dims.back() / 2; + CHECK_GT(hidden, 0); + output_dims.back() = hidden; + + auto output = std::make_shared(output_dims, input->Dtype(), input->GetDevice()); + const float *input_ptr = static_cast(input->DataPtr()); + float *output_ptr = static_cast(output->DataPtr()); + const int64_t rows = output->NumElements() / hidden; + for (int64_t row = 0; row < rows; ++row) { + const int64_t input_base = row * 2 * hidden; + const int64_t output_base = row * hidden; + for (int64_t col = 0; col < hidden; ++col) { + const float gate = input_ptr[input_base + col]; + const float up = input_ptr[input_base + hidden + col]; + output_ptr[output_base + col] = up * gate / (1.0f + std::exp(-gate)); + } + } + return output; +} + +std::shared_ptr SwiGLUBackward(const std::shared_ptr &input, + const std::shared_ptr &grad_output) { + CHECK(input->Dtype() == DataType::kFLOAT32); + CHECK(grad_output->Dtype() == input->Dtype()); + CHECK(input->IsContiguous()); + CHECK(grad_output->IsContiguous()); + CHECK_GT(input->Dims().size(), 0); + const int64_t hidden = input->Dims().back() / 2; + CHECK_GT(hidden, 0); + CHECK_EQ(grad_output->NumElements() * 2, input->NumElements()); + + auto grad_input = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + const float *input_ptr = static_cast(input->DataPtr()); + const float *grad_output_ptr = static_cast(grad_output->DataPtr()); + float *grad_input_ptr = static_cast(grad_input->DataPtr()); + const int64_t rows = grad_output->NumElements() / hidden; + for (int64_t row = 0; row < rows; ++row) { + const int64_t input_base = row * 2 * hidden; + const int64_t output_base = row * hidden; + for (int64_t col = 0; col < hidden; ++col) { + const float gate = input_ptr[input_base + col]; + const float up = input_ptr[input_base + hidden + col]; + const float grad = grad_output_ptr[output_base + col]; + const float sigmoid = 1.0f / (1.0f + std::exp(-gate)); + grad_input_ptr[input_base + col] = grad * up * sigmoid * (1.0f + gate * (1.0f - sigmoid)); + grad_input_ptr[input_base + hidden + col] = grad * gate * sigmoid; + } + } + return grad_input; +} +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_SWIGLU_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_SWIGLU_KERNEL(SwiGLUForward) +REGISTER_CPU_SWIGLU_KERNEL(SwiGLUBackward) + +#undef REGISTER_CPU_SWIGLU_KERNEL diff --git a/infini_train/src/kernels/cuda/elementwise.cu b/infini_train/src/kernels/cuda/elementwise.cu index fc423b35f..f2f8ba977 100644 --- a/infini_train/src/kernels/cuda/elementwise.cu +++ b/infini_train/src/kernels/cuda/elementwise.cu @@ -193,20 +193,6 @@ __global__ void BinaryBackwardKernelNoBroadcastVectorized(T *__restrict__ outA, } } -// Helper to choose optimal block size based on tensor size -inline size_t ChooseBlockSize(size_t num_elements) { - if (num_elements < 1024) { - return 64; - } - if (num_elements < 65536) { - return 128; - } - if (num_elements < 1048576) { - return 256; - } - return 512; -} - inline dim3 ChooseBlockDims(size_t num_elements) { return dim3(ChooseBlockSize(num_elements)); } // launch the given kernel function with the given output and inputs diff --git a/infini_train/src/kernels/cuda/swiglu.cu b/infini_train/src/kernels/cuda/swiglu.cu new file mode 100644 index 000000000..79ea92767 --- /dev/null +++ b/infini_train/src/kernels/cuda/swiglu.cu @@ -0,0 +1,137 @@ +#include +#include + +#include "infini_train/include/common/common.h" +#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/datatype.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/core/runtime/cuda/cuda_runtime_common.h" + +namespace infini_train::kernels::cuda { +namespace { +using namespace infini_train::common::cuda; + +// TODO(zbl): Optimize the packed [gate, up] accesses with vectorized loads/stores. +template +__global__ void SwiGLUForwardKernel(T *__restrict__ output, const T *__restrict__ input, int64_t hidden, + size_t num_elements) { + const size_t grid_stride = static_cast(gridDim.x) * blockDim.x; + for (size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < num_elements; + idx += grid_stride) { + const size_t row = idx / hidden; + const size_t col = idx % hidden; + const size_t input_base = row * 2 * hidden; + const T gate = input[input_base + col]; + const T up = input[input_base + hidden + col]; + output[idx] = Mul(up, Mul(gate, Sigmoid(gate))); + } +} + +template +__global__ void SwiGLUBackwardKernel(T *__restrict__ grad_input, const InputT *__restrict__ input, + const GradT *__restrict__ grad_output, int64_t hidden, size_t num_elements) { + const size_t grid_stride = static_cast(gridDim.x) * blockDim.x; + for (size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < num_elements; + idx += grid_stride) { + const size_t row = idx / hidden; + const size_t col = idx % hidden; + const size_t input_base = row * 2 * hidden; + const T gate = Cast(input[input_base + col]); + const T up = Cast(input[input_base + hidden + col]); + const T grad = Cast(grad_output[idx]); + const T sigmoid = Sigmoid(gate); + grad_input[input_base + col] = Mul(grad, Mul(up, Mul(sigmoid, Add(T(1), Mul(gate, Sub(T(1), sigmoid)))))); + grad_input[input_base + hidden + col] = Mul(grad, Mul(gate, sigmoid)); + } +} + +} // namespace + +std::shared_ptr SwiGLUForward(const std::shared_ptr &input) { + CHECK(input->IsContiguous()); + auto output_dims = input->Dims(); + CHECK_GT(output_dims.size(), 0); + CHECK_EQ(output_dims.back() % 2, 0); + const int64_t hidden = output_dims.back() / 2; + CHECK_GT(hidden, 0); + output_dims.back() = hidden; + + auto output = std::make_shared(output_dims, input->Dtype(), input->GetDevice()); + auto device = output->GetDevice(); + const auto &stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); + const size_t num_elements = output->NumElements(); + const dim3 block(ChooseBlockSize(num_elements)); + const dim3 grid(std::min(CEIL_DIV(num_elements, block.x), static_cast(65535))); + + switch (input->Dtype()) { + case DataType::kFLOAT32: + SwiGLUForwardKernel<<>>(static_cast(output->DataPtr()), + static_cast(input->DataPtr()), hidden, + num_elements); + break; + case DataType::kBFLOAT16: + SwiGLUForwardKernel<<>>(static_cast(output->DataPtr()), + static_cast(input->DataPtr()), hidden, + num_elements); + break; + default: + LOG_LOC(FATAL, "CUDA SwiGLUForward: unsupported data type"); + } + return output; +} + +std::shared_ptr SwiGLUBackward(const std::shared_ptr &input, + const std::shared_ptr &grad_output) { + CHECK(input->IsContiguous()); + CHECK(grad_output->IsContiguous()); + CHECK_GT(input->Dims().size(), 0); + const int64_t hidden = input->Dims().back() / 2; + CHECK_GT(hidden, 0); + CHECK_EQ(grad_output->NumElements() * 2, input->NumElements()); + + const DataType output_dtype = PromoteDataTypes(input->Dtype(), grad_output->Dtype()); + auto grad_input = std::make_shared(input->Dims(), output_dtype, input->GetDevice()); + auto device = input->GetDevice(); + const auto &stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); + const size_t num_elements = grad_output->NumElements(); + const dim3 block(ChooseBlockSize(num_elements)); + const dim3 grid(std::min(CEIL_DIV(num_elements, block.x), static_cast(65535))); + + if (input->Dtype() == DataType::kFLOAT32 && grad_output->Dtype() == DataType::kFLOAT32) { + SwiGLUBackwardKernel<<>>( + static_cast(grad_input->DataPtr()), static_cast(input->DataPtr()), + static_cast(grad_output->DataPtr()), hidden, num_elements); + } else if (input->Dtype() == DataType::kBFLOAT16 && grad_output->Dtype() == DataType::kBFLOAT16) { + SwiGLUBackwardKernel<<>>( + static_cast(grad_input->DataPtr()), static_cast(input->DataPtr()), + static_cast(grad_output->DataPtr()), hidden, num_elements); + } else if (input->Dtype() == DataType::kBFLOAT16 && grad_output->Dtype() == DataType::kFLOAT32) { + SwiGLUBackwardKernel<<>>( + static_cast(grad_input->DataPtr()), static_cast(input->DataPtr()), + static_cast(grad_output->DataPtr()), hidden, num_elements); + } else if (input->Dtype() == DataType::kFLOAT32 && grad_output->Dtype() == DataType::kBFLOAT16) { + SwiGLUBackwardKernel<<>>( + static_cast(grad_input->DataPtr()), static_cast(input->DataPtr()), + static_cast(grad_output->DataPtr()), hidden, num_elements); + } else { + LOG_LOC(FATAL, "CUDA SwiGLUBackward: unsupported data type combination"); + } + return grad_input; +} +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_SWIGLU_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_SWIGLU_KERNEL(SwiGLUForward) +REGISTER_CUDA_SWIGLU_KERNEL(SwiGLUBackward) + +#undef REGISTER_CUDA_SWIGLU_KERNEL diff --git a/infini_train/src/nn/lora/lora_utils.cc b/infini_train/src/nn/lora/lora_utils.cc index 10223d81d..45b6c243e 100644 --- a/infini_train/src/nn/lora/lora_utils.cc +++ b/infini_train/src/nn/lora/lora_utils.cc @@ -15,9 +15,11 @@ #include "infini_train/include/nn/functional.h" #include "infini_train/include/nn/lora/lora_linear.h" #include "infini_train/include/nn/lora/lora_parallel_linear.h" +#include "infini_train/include/nn/modules/activations.h" #include "infini_train/include/nn/modules/linear.h" #include "infini_train/include/nn/modules/module.h" #include "infini_train/include/nn/modules/transformer/causal_self_attention.h" +#include "infini_train/include/nn/modules/transformer/mlp.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" @@ -36,6 +38,9 @@ enum class LoRATensorSharding { // Attention QKV LoRA-B: dim0-sharded like ColumnParallel, but each local // shard is packed as [Qi | Ki | Vi] instead of a simple contiguous slice. kPackedQKVColumnParallelDim0, + // SwiGLU FC1 LoRA-B: dim0-sharded like ColumnParallel, but each local + // shard is packed as [gate_i | up_i] instead of a simple contiguous slice. + kPackedSwiGLUColumnParallelDim0, // LoRA-A for a RowParallelLinear: local shape is [rank, in/tp]. kRowParallelDim1, }; @@ -82,6 +87,41 @@ std::string QualifiedParamName(const std::string &module_name, const std::string return module_name.empty() ? param_name : module_name + "." + param_name; } +void MarkLoRAColumnParallelBSharding(const std::string &module_name, const std::shared_ptr &module, + const std::string &projection_name, LoRATensorSharding sharding, + std::unordered_map &shardings) { + auto projection = module->mutable_module(projection_name); + if (!dynamic_cast(projection.get())) { + return; + } + + const auto projection_module_name = QualifiedParamName(module_name, projection_name); + shardings[QualifiedParamName(projection_module_name, LoRAColumnParallelLinear::kParamLoraBName)] = sharding; +} + +void MarkPackedQKVLoRASharding(const std::string &module_name, const std::shared_ptr &module, + std::unordered_map &shardings) { + MarkLoRAColumnParallelBSharding(module_name, module, CausalSelfAttention::kCAttnLayerName, + LoRATensorSharding::kPackedQKVColumnParallelDim0, shardings); +} + +void MarkPackedSwiGLULoRASharding(const std::string &module_name, const std::shared_ptr &module, + std::unordered_map &shardings) { + bool is_swiglu = false; + for (const auto &child : module->modules()) { + if (dynamic_cast(child.get())) { + is_swiglu = true; + break; + } + } + if (!is_swiglu) { + return; + } + + MarkLoRAColumnParallelBSharding(module_name, module, MLP::kCFcLayerName, + LoRATensorSharding::kPackedSwiGLUColumnParallelDim0, shardings); +} + std::vector SortedLoRAStateDictNames(const std::unordered_map> &state_dict) { std::vector names; @@ -114,21 +154,14 @@ std::unordered_map BuildLoRATensorShardings(con } } - // Packed QKV is a property of the attention module topology, not the - // parameter name. Mark the LoRA-B of the attention QKV projection explicitly. + // Packed projections are properties of their parent module topology, not + // only their parameter names. for (const auto &[module_name, module] : named_modules) { - if (!dynamic_cast(module.get())) { - continue; + if (dynamic_cast(module.get())) { + MarkPackedQKVLoRASharding(module_name, module, shardings); + } else if (dynamic_cast(module.get())) { + MarkPackedSwiGLULoRASharding(module_name, module, shardings); } - - auto qkv_projection = module->mutable_module(CausalSelfAttention::kCAttnLayerName); - if (!dynamic_cast(qkv_projection.get())) { - continue; - } - - const auto qkv_module_name = QualifiedParamName(module_name, CausalSelfAttention::kCAttnLayerName); - shardings[QualifiedParamName(qkv_module_name, LoRAColumnParallelLinear::kParamLoraBName)] - = LoRATensorSharding::kPackedQKVColumnParallelDim0; } return shardings; @@ -165,14 +198,22 @@ ExportLoRATensorForSave(const std::string &name, const std::shared_ptr & case LoRATensorSharding::kReplicated: return tensor; case LoRATensorSharding::kPackedQKVColumnParallelDim0: - // Packed QKV is still dim0-sharded like ColumnParallel; it only needs - // an extra Q/K/V reorder after the common gather below. + case LoRATensorSharding::kPackedSwiGLUColumnParallelDim0: + // Packed projections are still dim0-sharded like ColumnParallel; they + // only need an extra component reorder after the common gather below. case LoRATensorSharding::kColumnParallelDim0: { auto gathered = parallel::GatherTensorParallelShard(tensor, 0); - if (sharding != LoRATensorSharding::kPackedQKVColumnParallelDim0) { + if (sharding == LoRATensorSharding::kColumnParallelDim0) { return gathered; } + if (sharding == LoRATensorSharding::kPackedSwiGLUColumnParallelDim0) { + // Local shards are gathered as [gate_0 | up_0 | gate_1 | up_1 | ...]. + // Adapter files use the full packed order [all_gate | all_up]. + return detail::RestorePackedSwiGLURowsFromTensorParallel(gathered, + parallel::global::GetTensorParallelSize()); + } + const auto lora_a_name = LoraANameForLoraB(name); auto lora_a_it = local_state_dict.find(lora_a_name); CHECK(lora_a_it != local_state_dict.end()) @@ -226,6 +267,14 @@ void LoadLoRATensorIntoModel(const std::string &name, const std::shared_ptrCopyFrom(sliced); return; } + if (sharding == LoRATensorSharding::kPackedSwiGLUColumnParallelDim0) { + // Full adapter files store [all_gate | all_up]. Each TP rank needs its + // matching local packed layout [gate_i | up_i]. + auto sliced = detail::SlicePackedSwiGLURowsForTensorParallel(src, parallel::tp_rank, tp_size); + CHECK(sliced->Dims() == dst_dims) << "LoadLoRATensorIntoModel: packed SwiGLU shard shape mismatch for " << name; + dst->CopyFrom(sliced); + return; + } CHECK_EQ(src_dims.size(), dst_dims.size()) << "LoadLoRATensorIntoModel: rank mismatch for " << name; int shard_dim = -1; @@ -246,82 +295,129 @@ void LoadLoRATensorIntoModel(const std::string &name, const std::shared_ptrCopyFrom(sliced); } -} // namespace - -namespace detail { - -// TODO: Reuse this packed-QKV sharding logic in TP checkpoint loading once the checkpoint infrastructure is stable. -// The current TP loader reads rank-local weights directly by file offset instead of slicing a materialized full tensor. -std::shared_ptr SlicePackedQKVRowsForTensorParallel(const std::shared_ptr &full_tensor, int64_t q_rows, - int tp_rank, int tp_size) { +std::shared_ptr SlicePackedProjectionRowsForTensorParallel(const std::shared_ptr &full_tensor, + const std::vector &projection_rows, + int tp_rank, int tp_size) { CHECK(full_tensor != nullptr); + CHECK(!projection_rows.empty()); const auto &dims = full_tensor->Dims(); CHECK_GE(dims.size(), 1); CHECK_GT(tp_size, 0); CHECK_GE(tp_rank, 0); CHECK_LT(tp_rank, tp_size); - CHECK_GT(q_rows, 0); - const int64_t total_rows = dims[0]; - CHECK_GT(total_rows, q_rows) << "Packed QKV tensor must contain Q, K, and V rows"; - CHECK_EQ((total_rows - q_rows) % 2, 0) << "Packed QKV K/V rows must be balanced"; + int64_t total_rows = 0; + for (int64_t rows : projection_rows) { + CHECK_GT(rows, 0); + CHECK_EQ(rows % tp_size, 0) << "Packed projection rows must be divisible by TP size"; + total_rows += rows; + } + CHECK_EQ(dims[0], total_rows) << "Packed projection row counts do not match tensor shape"; + + std::vector> shards; + shards.reserve(projection_rows.size()); + int64_t offset = 0; + for (int64_t rows : projection_rows) { + const int64_t local_rows = rows / tp_size; + const int64_t start = offset + static_cast(tp_rank) * local_rows; + shards.push_back(full_tensor->Slice(0, start, start + local_rows)); + offset += rows; + } + return nn::function::Concat(shards, 0); +} - const int64_t kv_rows = (total_rows - q_rows) / 2; - CHECK_GT(kv_rows, 0); - CHECK_EQ(q_rows % tp_size, 0) << "Q rows must be divisible by TP size"; - CHECK_EQ(kv_rows % tp_size, 0) << "K/V rows must be divisible by TP size"; +std::shared_ptr RestorePackedProjectionRowsFromTensorParallel(const std::shared_ptr &gathered_tensor, + const std::vector &local_projection_rows, + int tp_size) { + CHECK(gathered_tensor != nullptr); + CHECK(!local_projection_rows.empty()); - const int64_t q_local_rows = q_rows / tp_size; - const int64_t kv_local_rows = kv_rows / tp_size; - CHECK_GT(q_local_rows, 0); - CHECK_GT(kv_local_rows, 0); + const auto &dims = gathered_tensor->Dims(); + CHECK_GE(dims.size(), 1); + CHECK_GT(tp_size, 0); + CHECK_EQ(dims[0] % tp_size, 0) << "Gathered packed projection rows must be divisible by TP size"; + + int64_t rows_per_rank = 0; + for (int64_t rows : local_projection_rows) { + CHECK_GT(rows, 0); + rows_per_rank += rows; + } + CHECK_EQ(dims[0] / tp_size, rows_per_rank) + << "Local packed projection row counts do not match gathered tensor shape"; + + std::vector> reordered_shards; + reordered_shards.reserve(static_cast(tp_size) * local_projection_rows.size()); + int64_t projection_offset = 0; + for (int64_t local_rows : local_projection_rows) { + for (int rank = 0; rank < tp_size; ++rank) { + const int64_t base = static_cast(rank) * rows_per_rank + projection_offset; + reordered_shards.push_back(gathered_tensor->Slice(0, base, base + local_rows)); + } + projection_offset += local_rows; + } + return nn::function::Concat(reordered_shards, 0); +} + +} // namespace - auto q_shard = full_tensor->Slice(0, static_cast(tp_rank) * q_local_rows, - static_cast(tp_rank + 1) * q_local_rows); - auto k_shard = full_tensor->Slice(0, q_rows + static_cast(tp_rank) * kv_local_rows, - q_rows + static_cast(tp_rank + 1) * kv_local_rows); - auto v_shard = full_tensor->Slice(0, q_rows + kv_rows + static_cast(tp_rank) * kv_local_rows, - q_rows + kv_rows + static_cast(tp_rank + 1) * kv_local_rows); +namespace detail { + +// TODO: Reuse this packed-QKV sharding logic in TP checkpoint loading once the checkpoint infrastructure is stable. +// The current TP loader reads rank-local weights directly by file offset instead of slicing a materialized full tensor. +std::shared_ptr SlicePackedQKVRowsForTensorParallel(const std::shared_ptr &full_tensor, int64_t q_rows, + int tp_rank, int tp_size) { + CHECK(full_tensor != nullptr); + CHECK_GE(full_tensor->Dims().size(), 1); + CHECK_GT(q_rows, 0); + CHECK_GT(full_tensor->Dims()[0], q_rows) << "Packed QKV tensor must contain Q, K, and V rows"; + CHECK_EQ((full_tensor->Dims()[0] - q_rows) % 2, 0) << "Packed QKV K/V rows must be balanced"; - return infini_train::nn::function::Concat({q_shard, k_shard, v_shard}, 0); + const int64_t kv_rows = (full_tensor->Dims()[0] - q_rows) / 2; + CHECK_GT(kv_rows, 0); + return SlicePackedProjectionRowsForTensorParallel(full_tensor, {q_rows, kv_rows, kv_rows}, tp_rank, tp_size); } std::shared_ptr RestorePackedQKVRowsFromTensorParallel(const std::shared_ptr &gathered_tensor, int64_t q_rows, int tp_size) { CHECK(gathered_tensor != nullptr); - - const auto &dims = gathered_tensor->Dims(); - CHECK_GE(dims.size(), 1); + CHECK_GE(gathered_tensor->Dims().size(), 1); CHECK_GT(tp_size, 0); CHECK_GT(q_rows, 0); - CHECK_EQ(dims[0] % tp_size, 0) << "Gathered packed QKV rows must be divisible by TP size"; - - const int64_t local_rows = dims[0] / tp_size; CHECK_EQ(q_rows % tp_size, 0) << "Q rows must be divisible by TP size"; + const int64_t q_local_rows = q_rows / tp_size; + const int64_t local_rows = gathered_tensor->Dims()[0] / tp_size; CHECK_GT(local_rows, q_local_rows) << "Gathered packed QKV tensor must contain local Q, K, and V rows"; CHECK_EQ((local_rows - q_local_rows) % 2, 0) << "Local packed QKV K/V rows must be balanced"; const int64_t kv_local_rows = (local_rows - q_local_rows) / 2; CHECK_GT(kv_local_rows, 0); + return RestorePackedProjectionRowsFromTensorParallel(gathered_tensor, {q_local_rows, kv_local_rows, kv_local_rows}, + tp_size); +} - std::vector> reordered_shards; - reordered_shards.reserve(static_cast(tp_size) * 3); - for (int rank = 0; rank < tp_size; ++rank) { - const int64_t base = static_cast(rank) * local_rows; - reordered_shards.push_back(gathered_tensor->Slice(0, base, base + q_local_rows)); - } - for (int rank = 0; rank < tp_size; ++rank) { - const int64_t base = static_cast(rank) * local_rows; - reordered_shards.push_back(gathered_tensor->Slice(0, base + q_local_rows, base + q_local_rows + kv_local_rows)); - } - for (int rank = 0; rank < tp_size; ++rank) { - const int64_t base = static_cast(rank) * local_rows; - reordered_shards.push_back( - gathered_tensor->Slice(0, base + q_local_rows + kv_local_rows, base + q_local_rows + 2 * kv_local_rows)); - } +std::shared_ptr SlicePackedSwiGLURowsForTensorParallel(const std::shared_ptr &full_tensor, int tp_rank, + int tp_size) { + CHECK(full_tensor != nullptr); + CHECK_GE(full_tensor->Dims().size(), 1); + CHECK_EQ(full_tensor->Dims()[0] % 2, 0) << "Packed SwiGLU tensor must contain balanced gate and up rows"; - return nn::function::Concat(reordered_shards, 0); + const int64_t rows_per_projection = full_tensor->Dims()[0] / 2; + CHECK_GT(rows_per_projection, 0); + return SlicePackedProjectionRowsForTensorParallel(full_tensor, {rows_per_projection, rows_per_projection}, tp_rank, + tp_size); +} + +std::shared_ptr RestorePackedSwiGLURowsFromTensorParallel(const std::shared_ptr &gathered_tensor, + int tp_size) { + CHECK(gathered_tensor != nullptr); + CHECK_GE(gathered_tensor->Dims().size(), 1); + CHECK_GT(tp_size, 0); + const int64_t rows_per_rank = gathered_tensor->Dims()[0] / tp_size; + CHECK_EQ(rows_per_rank % 2, 0) << "Local packed SwiGLU tensor must contain balanced gate and up rows"; + const int64_t local_rows = rows_per_rank / 2; + CHECK_GT(local_rows, 0); + return RestorePackedProjectionRowsFromTensorParallel(gathered_tensor, {local_rows, local_rows}, tp_size); } } // namespace detail diff --git a/infini_train/src/nn/modules/activations.cc b/infini_train/src/nn/modules/activations.cc index d1bbc9da8..37a578059 100644 --- a/infini_train/src/nn/modules/activations.cc +++ b/infini_train/src/nn/modules/activations.cc @@ -19,6 +19,6 @@ std::vector> NewGELU::Forward(const std::vector> SwiGLU::Forward(const std::vector> &x) { - return {x[0] * function::Sigmoid(x[0])}; + return std::make_shared()->Apply(x); } } // namespace infini_train::nn diff --git a/infini_train/src/nn/modules/transformer/mlp.cc b/infini_train/src/nn/modules/transformer/mlp.cc index 115d77e6f..3cf94f91c 100644 --- a/infini_train/src/nn/modules/transformer/mlp.cc +++ b/infini_train/src/nn/modules/transformer/mlp.cc @@ -48,29 +48,19 @@ MLP::MLP(const TransformerConfig &config) : CloneableModule(kType) { // c_fc: ColumnParallel (input full, output parallel) modules_[kCFcLayerName] = std::make_shared( - /*in_features=*/config.n_embd, /*out_features=*/ffn_hidden, + /*in_features=*/config.n_embd, + /*out_features=*/config.activation_type == MLPType::kSwiGLU ? 2 * ffn_hidden : ffn_hidden, /*bias=*/config.add_bias_linear, /*gather_output=*/false, /*input_is_parallel=*/false, /*skip_bias_add=*/false, /*sequence_parallel=*/parallel::global::GetSequenceParallelEnabled()); - // For SwiGLU, add second projection - if (config.activation_type == MLPType::kSwiGLU) { - modules_[kCFc2LayerName] = std::make_shared( - /*in_features=*/config.n_embd, /*out_features=*/ffn_hidden, - /*bias=*/config.add_bias_linear, - /*gather_output=*/false, - /*input_is_parallel=*/false, - /*skip_bias_add=*/false, - /*sequence_parallel=*/nn::parallel::global::GetSequenceParallelEnabled()); - } - // Activation: check for GELU or SwiGLU if (config.activation_type == MLPType::kGELU) { modules_[kGeluLayerName] = std::make_shared(); } else if (config.activation_type == MLPType::kSwiGLU) { - modules_[kSiluLayerName] = std::make_shared(); + modules_[kSwiGLULayerName] = std::make_shared(); } // c_proj: RowParallel (input parallel, output full) @@ -85,31 +75,22 @@ MLP::MLP(const TransformerConfig &config) : CloneableModule(kType) { std::vector> MLP::Forward(const std::vector> &x) { - bool is_swiglu = modules_.contains(kCFc2LayerName) && modules_.contains(kSiluLayerName); - - if (is_swiglu) { - // SwiGLU forward pass - // (B, T, C) -> ColumnParallelLinear(C, hidden_dim) -> (B, T, hidden_dim) - auto x1 = (*modules_[kCFcLayerName])(x)[0]; - // (B, T, C) -> ColumnParallelLinear(C, hidden_dim) -> (B, T, hidden_dim) - auto x2 = (*modules_[kCFc2LayerName])(x)[0]; - // (B, T, hidden_dim) -> SiLU -> (B, T, hidden_dim) - x2 = (*modules_[kSiluLayerName])({x2})[0]; - // (B, T, hidden_dim) -> element-wise mul -> (B, T, hidden_dim) - auto x3 = x1 * x2; - // (B, T, hidden_dim) -> RowParallelLinear(hidden_dim, C) -> (B, T, C) - auto x4 = (*modules_[kCProjLayerName])({x3}); - return x4; - } else { - // GELU forward pass (standard) - // (B, T, C) -> ColumnParallelLinear(C, 4*C) -> (B, T, 4*C_local) - auto x1 = (*modules_[kCFcLayerName])(x); - // (B, T, 4*C_local) -> GELU -> (B, T, 4*C_local) - auto x2 = (*modules_[kGeluLayerName])(x1); - // (B, T, 4*C_local) -> RowParallelLinear(4*C, C) -> (B, T, C) - auto x3 = (*modules_[kCProjLayerName])(x2); - return x3; + if (modules_.contains(kSwiGLULayerName)) { + // (B, T, C) -> ColumnParallelLinear(C, 2*H) -> (B, T, 2*H_local) + auto packed = (*modules_[kCFcLayerName])(x)[0]; + // (B, T, 2*H_local) [gate, up] -> SwiGLU -> (B, T, H_local) + auto activated = (*modules_[kSwiGLULayerName])({packed}); + // (B, T, H_local) -> RowParallelLinear(H, C) -> (B, T, C) + return (*modules_[kCProjLayerName])(activated); } + + // GELU forward pass (standard) + // (B, T, C) -> ColumnParallelLinear(C, 4*C) -> (B, T, 4*C_local) + auto x1 = (*modules_[kCFcLayerName])(x); + // (B, T, 4*C_local) -> GELU -> (B, T, 4*C_local) + auto x2 = (*modules_[kGeluLayerName])(x1); + // (B, T, 4*C_local) -> RowParallelLinear(4*C, C) -> (B, T, C) + return (*modules_[kCProjLayerName])(x2); } } // namespace infini_train::nn diff --git a/tests/autograd/test_autograd_elementwise_backward.cc b/tests/autograd/test_autograd_elementwise_backward.cc index f7eb0d5ff..fbd842d8f 100644 --- a/tests/autograd/test_autograd_elementwise_backward.cc +++ b/tests/autograd/test_autograd_elementwise_backward.cc @@ -3,6 +3,7 @@ #include "gtest/gtest.h" +#include "infini_train/include/autograd/activations.h" #include "infini_train/include/autograd/elementwise.h" #include "infini_train/include/nn/parallel/global.h" #include "infini_train/include/tensor.h" @@ -13,6 +14,70 @@ using namespace infini_train; class AutogradElementwiseBackwardTest : public infini_train::test::InfiniTrainTest {}; +TEST_P(AutogradElementwiseBackwardTest, SwiGLUForwardBackward) { + const std::vector input_dims{2, 6}; + const std::vector input_values{-1.0f, 0.5f, 2.0f, -2.0f, 0.0f, 1.5f, 0.25f, -3.0f, 1.0f, 0.5f, -1.0f, 2.0f}; + const std::vector grad_values{1.0f, -0.5f, 2.0f, -1.5f, 0.25f, 0.75f}; + auto input = std::make_shared(input_values.data(), input_dims, DataType::kFLOAT32, GetDevice()); + auto grad_output + = std::make_shared(grad_values.data(), std::vector{2, 3}, DataType::kFLOAT32, GetDevice()); + + std::vector expected_output(6); + std::vector expected_grad(12); + for (int64_t row = 0; row < 2; ++row) { + for (int64_t col = 0; col < 3; ++col) { + const int64_t packed_base = row * 6; + const int64_t output_idx = row * 3 + col; + const float gate = input_values[packed_base + col]; + const float up = input_values[packed_base + 3 + col]; + const float grad = grad_values[output_idx]; + const float sigmoid = 1.0f / (1.0f + std::exp(-gate)); + expected_output[output_idx] = up * gate * sigmoid; + expected_grad[packed_base + col] = grad * up * sigmoid * (1.0f + gate * (1.0f - sigmoid)); + expected_grad[packed_base + 3 + col] = grad * gate * sigmoid; + } + } + + auto swiglu_fn = std::make_shared(); + auto result = swiglu_fn->Apply({input}); + ASSERT_EQ(result.size(), 1); + EXPECT_EQ(result[0]->Dims(), (std::vector{2, 3})); + test::ExpectTensorNear(result[0], expected_output, 1e-5f); + + auto grad_inputs = swiglu_fn->Backward({grad_output}); + ASSERT_EQ(grad_inputs.size(), 1); + EXPECT_EQ(grad_inputs[0]->Dims(), input_dims); + test::ExpectTensorNear(grad_inputs[0], expected_grad, 1e-5f); +} + +TEST_P(AutogradElementwiseBackwardTest, SwiGLUAutocastBackward) { + SKIP_CPU(); + const std::vector input_dims{1, 4}; + const std::vector input_values{0.5f, -1.0f, 1.0f, -0.5f}; + const std::vector grad_values{2.0f, -0.25f}; + auto input_fp32 = std::make_shared(input_values.data(), input_dims, DataType::kFLOAT32, GetDevice()); + auto input = std::make_shared(input_fp32->To(DataType::kBFLOAT16)); + auto grad_output + = std::make_shared(grad_values.data(), std::vector{1, 2}, DataType::kFLOAT32, GetDevice()); + + auto swiglu_fn = std::make_shared(); + swiglu_fn->Apply({input}); + auto grad_inputs = swiglu_fn->Backward({grad_output}); + ASSERT_EQ(grad_inputs.size(), 1); + EXPECT_EQ(grad_inputs[0]->Dtype(), DataType::kFLOAT32); + + std::vector expected_grad(4); + for (int64_t col = 0; col < 2; ++col) { + const float gate = input_values[col]; + const float up = input_values[2 + col]; + const float grad = grad_values[col]; + const float sigmoid = 1.0f / (1.0f + std::exp(-gate)); + expected_grad[col] = grad * up * sigmoid * (1.0f + gate * (1.0f - sigmoid)); + expected_grad[2 + col] = grad * gate * sigmoid; + } + test::ExpectTensorNear(grad_inputs[0], expected_grad, 2e-3f); +} + TEST_P(AutogradElementwiseBackwardTest, AddBackward) { auto a = std::make_shared(std::vector{2, 3}, DataType::kFLOAT32, GetDevice(), true); a->Fill(1.0f); diff --git a/tests/lora/test_lora.cc b/tests/lora/test_lora.cc index 831a59df6..26cffdcaa 100644 --- a/tests/lora/test_lora.cc +++ b/tests/lora/test_lora.cc @@ -105,6 +105,24 @@ TEST_P(LoRATest, PackedQKVRestoreFromTPGather) { ExpectRows(restored, {0, 1, 6, 7, 2, 3, 8, 9, 4, 5, 10, 11}); } +TEST_P(LoRATest, PackedSwiGLUShard) { + auto full_swiglu = MakeRowLabeledTensor(/*rows=*/8, /*cols=*/3, GetDevice()); + auto shard = infini_train::nn::lora::detail::SlicePackedSwiGLURowsForTensorParallel(full_swiglu, /*tp_rank=*/1, + /*tp_size=*/2); + + EXPECT_EQ(shard->Dims(), (std::vector{4, 3})); + ExpectRows(shard, {2, 3, 6, 7}); +} + +TEST_P(LoRATest, PackedSwiGLURestoreFromTPGather) { + auto rank_major_swiglu = MakeRowLabeledTensor(/*rows=*/8, /*cols=*/3, GetDevice()); + auto restored = infini_train::nn::lora::detail::RestorePackedSwiGLURowsFromTensorParallel(rank_major_swiglu, + /*tp_size=*/2); + + EXPECT_EQ(restored->Dims(), (std::vector{8, 3})); + ExpectRows(restored, {0, 1, 4, 5, 2, 3, 6, 7}); +} + TEST_P(LoRATest, LoRAConfigShouldApply) { LoRAConfig config; config.rank = 8; diff --git a/tests/transformer/test_transformer_architecture.cc b/tests/transformer/test_transformer_architecture.cc index d4a6efc29..4cec471de 100644 --- a/tests/transformer/test_transformer_architecture.cc +++ b/tests/transformer/test_transformer_architecture.cc @@ -88,7 +88,11 @@ TEST_P(TransformerModuleTest, SwiGLUMLP) { auto mlp = std::make_shared(config); mlp->To(GetDevice()); - EXPECT_EQ(mlp->Parameters().size(), 3); + EXPECT_EQ(mlp->Parameters().size(), 2); + auto state = mlp->StateDict(); + ASSERT_TRUE(state.contains("c_fc.weight")); + EXPECT_FALSE(state.contains("c_fc2.weight")); + EXPECT_EQ(state.at("c_fc.weight")->Dims(), (std::vector{512, config.n_embd})); auto input = std::make_shared(std::vector{2, 8, 64}, DataType::kFLOAT32, GetDevice()); auto output = (*mlp)({input}); @@ -241,10 +245,9 @@ TEST_P(TransformerModuleTest, MoELayerTop2SwiGLU) { auto state = moe->StateDict(); ASSERT_TRUE(state.contains("experts.expert_0.c_fc.weight")); - ASSERT_TRUE(state.contains("experts.expert_0.c_fc2.weight")); + EXPECT_FALSE(state.contains("experts.expert_0.c_fc2.weight")); ASSERT_TRUE(state.contains("experts.expert_0.c_proj.weight")); - EXPECT_EQ(state.at("experts.expert_0.c_fc.weight")->Dims(), (std::vector{48, config.n_embd})); - EXPECT_EQ(state.at("experts.expert_0.c_fc2.weight")->Dims(), (std::vector{48, config.n_embd})); + EXPECT_EQ(state.at("experts.expert_0.c_fc.weight")->Dims(), (std::vector{96, config.n_embd})); EXPECT_EQ(state.at("experts.expert_0.c_proj.weight")->Dims(), (std::vector{config.n_embd, 48})); }