diff --git a/include/flexflow/ffconst.h b/include/flexflow/ffconst.h index 7925a18a6..328e0a1ea 100644 --- a/include/flexflow/ffconst.h +++ b/include/flexflow/ffconst.h @@ -198,6 +198,7 @@ enum OperatorType { OP_RESIDUAL_RMS_NORM, OP_BEAM_TOPK, OP_ARGMAX, + OP_DECODING, OP_INC_MULTIHEAD_SELF_ATTENTION, OP_SPEC_INC_MULTIHEAD_SELF_ATTENTION, OP_TREE_INC_MULTIHEAD_SELF_ATTENTION, diff --git a/include/flexflow/model.h b/include/flexflow/model.h index 874211172..9a5f6dbca 100644 --- a/include/flexflow/model.h +++ b/include/flexflow/model.h @@ -168,6 +168,10 @@ enum TaskIDs { ARGMAX_INIT_TASK_ID, ARGMAX_BEAM_INF_TASK_ID, ARGMAX_NORM_INF_TASK_ID, + DECODING_INIT_TASK_ID, + DECODING_BEAM_INF_TASK_ID, + DECODING_NORM_INF_TASK_ID, + DECODING_PEFT_BWD_TASK_ID, TRANSPOSE_INIT_TASK_ID, TRANSPOSE_FWD_TASK_ID, TRANSPOSE_BWD_TASK_ID, @@ -375,6 +379,7 @@ class BeamTopK; class SpecIncMultiHeadSelfAttention; class Sampling; class ArgMax; +class Decoding; class Combine; class Repartition; class Reduction; @@ -720,6 +725,7 @@ class FFModel { bool speculative_decoding, char const *name = NULL); Tensor argmax(const Tensor input, bool beam_search, char const *name = NULL); + Tensor decoding(const Tensor input, bool beam_search, char const *name = NULL); Tensor sampling(const Tensor input, float top_p, char const *name = NULL); Tensor multihead_attention(const Tensor query, const Tensor key, @@ -1221,6 +1227,8 @@ class FFModel { Sampling *>, std::unordered_map, ArgMax *>, + std::unordered_map, + Decoding *>, std::unordered_map< std::pair, SpecIncMultiHeadSelfAttention *>, diff --git a/include/flexflow/operator_params.h b/include/flexflow/operator_params.h index 673f78ad4..b5c121ca7 100644 --- a/include/flexflow/operator_params.h +++ b/include/flexflow/operator_params.h @@ -6,6 +6,7 @@ #include "flexflow/ops/aggregate_spec_params.h" #include "flexflow/ops/arg_topk_params.h" #include "flexflow/ops/argmax_params.h" +#include "flexflow/ops/decoding_params.h" #include "flexflow/ops/attention_params.h" #include "flexflow/ops/batch_matmul_params.h" #include "flexflow/ops/beam_topk_params.h" @@ -85,6 +86,7 @@ using OperatorParameters = mp::variant const &, + std::vector const &, + MachineView const *mv = nullptr) override; + void forward(FFModel const &) override; + Legion::FutureMap inference(FFModel const &, + BatchConfigFuture const &, + std::vector const &, + std::vector const &, + MachineView const *mv = nullptr) override; + Legion::FutureMap peft_bwd(FFModel const &, + BatchConfigFuture const &, + std::vector const &, + std::vector const &, + MachineView const *mv = nullptr) override; + void backward(FFModel const &) override; + void print_layer(FFModel const &model) override { + assert(0); + } + static Op * + create_operator_from_layer(FFModel &model, + Layer const *layer, + std::vector const &inputs); + static OpMeta *init_task(Legion::Task const *task, + std::vector const ®ions, + Legion::Context ctx, + Legion::Runtime *runtime); + static BeamInferenceResult + inference_task_beam(Legion::Task const *task, + std::vector const ®ions, + Legion::Context ctx, + Legion::Runtime *runtime); + static InferenceResult + inference_task_norm(Legion::Task const *task, + std::vector const ®ions, + Legion::Context ctx, + Legion::Runtime *runtime); + static bool peft_bwd_task(Legion::Task const *task, + std::vector const ®ions, + Legion::Context ctx, + Legion::Runtime *runtime); + bool measure_operator_cost(Simulator *sim, + MachineView const &pc, + CostMetrics &cost_metrics) const override; + void serialize(Legion::Serializer &) const override; + static PCG::Node deserialize(FFModel &ff, + Legion::Deserializer &d, + ParallelTensor inputs[], + int num_inputs); + Op *materialize(FFModel &ff, + ParallelTensor inputs[], + int num_inputs) const override; + Params get_params() const; + + template + static void inference_kernel(DecodingMeta const *m, + BatchConfig const *bc, + DT const *input_ptr, + DT *softmax_output_ptr, + int *argmax_output_ptr, + int num_classes, + int vocab_offset, + float *loss, + ffStream_t stream); + static void inference_kernel_wrapper(DecodingMeta *m, + BatchConfig const *bc, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &softmax_output, + GenericTensorAccessorW const &argmax_output); + template + static void peft_bwd_kernel(DecodingMeta const *m, + BatchConfig const *bc, + DT *input_grad_ptr, + int num_classes, + int shard_id, + ffStream_t stream); + static void peft_bwd_kernel_wrapper(DecodingMeta *m, + BatchConfig const *bc, + int shard_id, + GenericTensorAccessorW const &input_grad); + +public: + LayerID layer_guid; + bool beam_search; +}; + +class DecodingMeta : public OpMeta { +public: + DecodingMeta(FFHandler handler, + Decoding const *decoding, + Legion::Domain const &input_domain, + MemoryAllocator &gpu_mem_allocator); + ~DecodingMeta(void); + bool beam_search; + float *probs; + float *d_loss; + // Temporary buffers + int *parent_output_buffer; + // Sharded softmax context + SoftmaxShardedContext *softmax_context; + // PEFT related fields + void *output_grad_ptr = nullptr; + size_t allocated_peft_buffer_size = 0; + Realm::RegionInstance reserveInst; + BatchConfig::TokenId peft_token_ids[BatchConfig::MAX_NUM_TOKENS]; +}; + +}; // namespace FlexFlow + +#endif // _FLEXFLOW_DECODING_H diff --git a/include/flexflow/ops/decoding_params.h b/include/flexflow/ops/decoding_params.h new file mode 100644 index 000000000..bee4a3174 --- /dev/null +++ b/include/flexflow/ops/decoding_params.h @@ -0,0 +1,26 @@ +#ifndef _FLEXFLOW_DECODING_PARAMS_H +#define _FLEXFLOW_DECODING_PARAMS_H + +#include "flexflow/ffconst.h" +#include "flexflow/parallel_tensor.h" + +namespace FlexFlow { + +struct DecodingParams { + LayerID layer_guid; + bool beam_search; + bool is_valid(ParallelTensorShape const &) const; + char name[MAX_OPNAME]; +}; +bool operator==(DecodingParams const &, DecodingParams const &); + +} // namespace FlexFlow + +namespace std { +template <> +struct hash { + size_t operator()(FlexFlow::DecodingParams const &) const; +}; +} // namespace std + +#endif // _FLEXFLOW_DECODING_PARAMS_H diff --git a/inference/models/llama.cc b/inference/models/llama.cc index 76abba0f5..86673d81a 100644 --- a/inference/models/llama.cc +++ b/inference/models/llama.cc @@ -291,8 +291,9 @@ void LLAMA::create_llama_model(FFModel &ff, output = ff.sampling(softmax, generation_config.topp); } else { // output = ff.arg_top_k(dense, /*k=*/1, false); - Tensor softmax = ff.softmax(dense, -1); - output = ff.argmax(softmax, /*beam_Search*/ false); + // Tensor softmax = ff.softmax(dense, -1); + // output = ff.argmax(softmax, /*beam_Search*/ false); + output = ff.decoding(dense, /*beam_search*/ false, "decoding"); } } diff --git a/src/ops/argmax.cc b/src/ops/argmax.cc index e953430b7..bd3bd797e 100644 --- a/src/ops/argmax.cc +++ b/src/ops/argmax.cc @@ -107,6 +107,22 @@ bool operator==(ArgMaxParams const &lhs, ArgMaxParams const &rhs) { return lhs.beam_search == rhs.beam_search; } +static std::string remove_uid(char const *op_name) { + std::string op_name_without_uid = std::string(op_name); + size_t last_underscore = op_name_without_uid.length(); + for (int i = op_name_without_uid.length() - 1; i > 0; i--) { + if (!(std::isdigit(op_name[i]) || op_name[i] == '_')) { + break; + } else if (op_name[i] == '_') { + last_underscore = i; + } + } + if (last_underscore < op_name_without_uid.length()) { + op_name_without_uid.erase(last_underscore); + } + return op_name_without_uid; +} + ArgMax::ArgMax(FFModel &model, const ParallelTensor _input, bool _beam_search, @@ -136,6 +152,10 @@ ArgMax::ArgMax(FFModel &model, outputs[1] = model.create_parallel_tensor_legion_ordering( numdim, dims, DT_INT32, this, 1 /*owner_idx*/); } + std::string const &input_label = std::string("Argmax input tensor"); + _input->print(input_label); + std::string const &label = std::string("Argmax output tensor"); + outputs[0]->print(label); } ArgMax::ArgMax(FFModel &model, ArgMax const &other, const ParallelTensor input) @@ -397,6 +417,18 @@ InferenceResult ArgMax::inference_kernel_wrapper(m, bc, input, indices, parent, &loss); + if (task->index_point.point_data[0] == 0) { + int in_dim0 = input.domain.hi()[0] - input.domain.lo()[0] + 1; + int in_dim1 = input.domain.hi()[1] - input.domain.lo()[1] + 1; + int out_dim0 = indices.domain.hi()[0] - indices.domain.lo()[0] + 1; + int out_dim1 = indices.domain.hi()[1] - indices.domain.lo()[1] + 1; + std::string op_name_without_uid = remove_uid(m->op_name); + printf("Argmax(%s): in=[%i, bz=%i/%i] -> out=[%i,bz=%i/%i]\n", + op_name_without_uid.c_str(), + in_dim0, bc->num_tokens, in_dim1, + out_dim0, bc->num_tokens, out_dim1); + } + InferenceResult ir; ir.finetuning_loss = loss; diff --git a/src/ops/decoding.cc b/src/ops/decoding.cc new file mode 100644 index 000000000..f85dc4416 --- /dev/null +++ b/src/ops/decoding.cc @@ -0,0 +1,621 @@ +/* Copyright 2023 CMU, Facebook, LANL, MIT, NVIDIA, and Stanford (alphabetical) + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "flexflow/ops/decoding.h" +#include "flexflow/model.h" +#include "flexflow/ops/fused.h" +#include "flexflow/utils/hash_utils.h" +#include "legion/legion_utilities.h" + +namespace FlexFlow { +// declare Legion names +using Legion::ArgumentMap; +using Legion::Context; +using Legion::coord_t; +using Legion::Domain; +using Legion::FutureMap; +using Legion::IndexLauncher; +using Legion::InlineLauncher; +using Legion::Machine; +using Legion::Memory; +using Legion::PhysicalRegion; +using Legion::Predicate; +using Legion::Rect; +using Legion::RegionRequirement; +using Legion::Runtime; +using Legion::Task; +using Legion::TaskArgument; +using Legion::TaskLauncher; +using PCG::Node; + + +/* Params */ +bool operator==(DecodingParams const &lhs, DecodingParams const &rhs) { + return lhs.layer_guid == rhs.layer_guid && lhs.beam_search == rhs.beam_search; +} + +void Decoding::serialize(Legion::Serializer &sez) const { + sez.serialize(this->layer_guid.id); + sez.serialize(this->layer_guid.transformer_layer_id); + sez.serialize(this->layer_guid.model_id); + sez.serialize(this->beam_search); + sez.serialize(strlen(this->name)); + sez.serialize(this->name, strlen(this->name)); +} + +/*static*/ +Node Decoding::deserialize(FFModel &ff, + Legion::Deserializer &dez, + ParallelTensor inputs[], + int num_inputs) { + assert(num_inputs == 1); + size_t id, transformer_layer_id, deserialized_model_id; + dez.deserialize(id); + dez.deserialize(transformer_layer_id); + dez.deserialize(deserialized_model_id); + LayerID layer_guid(id, transformer_layer_id, deserialized_model_id); + bool beam_search; + dez.deserialize(beam_search); + size_t name_len; + char name[MAX_OPNAME] = {0}; + dez.deserialize(name_len); + dez.deserialize(name, name_len); + + DecodingParams params; + params.layer_guid = layer_guid; + params.beam_search = beam_search; + strcpy(params.name, name); + return ff.get_or_create_node(inputs[0], params); +} + +bool DecodingParams::is_valid(ParallelTensorShape const &input) const { + return input.is_valid(); +} + +DecodingParams Decoding::get_params() const { + DecodingParams params; + params.layer_guid = this->layer_guid; + params.beam_search = this->beam_search; + if (strlen(this->name) < MAX_OPNAME) { + strcpy(params.name, this->name); + } + return params; +} + +Tensor FFModel::decoding(const Tensor input, bool beam_search, char const *name) { + Layer *li = new Layer(this, + OP_DECODING, + input->data_type, + name, + 1 /*inputs*/, + 0 /*weights*/, + 2 /*outputs*/, + input); + + printf("Adding decoding layer\n"); + { + int numdims = input->num_dims; + int dims[MAX_TENSOR_DIM]; + + // First output: softmax output (same dimensions as input) + for (int i = 0; i < numdims; i++) { + dims[i] = input->dims[i]; + } + li->outputs[0] = create_tensor_legion_ordering( + numdims, dims, input->data_type, li, 0, true /*create_grad*/); + + // Second output: argmax output (batch dimensions only, int type) + for (int i = 1; i < numdims; i++) { + dims[i-1] = input->dims[i]; // Shift batch dimensions down + } + int argmax_numdims = numdims - 1; // Remove vocab dimension + li->outputs[1] = create_tensor_legion_ordering( + argmax_numdims, dims, DT_INT32, li, 1, false /*create_grad*/); + } + li->add_int_property("beam_search", beam_search); + layers.push_back(li); + return li->outputs[1]; // Return argmax output for compatibility +} + +Op *Decoding::create_operator_from_layer( + FFModel &model, + Layer const *layer, + std::vector const &inputs) { + long long value; + layer->get_int_property("beam_search", value); + bool beam_search = (bool)value; + return new Decoding(model, layer->layer_guid, inputs[0], beam_search, layer->name); +} + +static std::string remove_uid(char const *op_name) { + std::string op_name_without_uid = std::string(op_name); + size_t last_underscore = op_name_without_uid.length(); + for (int i = op_name_without_uid.length() - 1; i > 0; i--) { + if (!(std::isdigit(op_name[i]) || op_name[i] == '_')) { + break; + } else if (op_name[i] == '_') { + last_underscore = i; + } + } + if (last_underscore < op_name_without_uid.length()) { + op_name_without_uid.erase(last_underscore); + } + return op_name_without_uid; +} + +Decoding::Decoding(FFModel &model, + LayerID const &_layer_guid, + const ParallelTensor _input, + bool _beam_search, + char const *name) + : Op(model, + OP_DECODING, + _input->data_type, + name, + 1 /*inputs*/, + 0 /*weights*/, + 2 /*outputs*/, + _input), + beam_search(_beam_search) { + layer_guid = _layer_guid; + int numdim = inputs[0]->num_dims; + ParallelDim dims[MAX_TENSOR_DIM]; + + // First output: softmax output (same dimensions as input) + for (int i = 0; i < numdim; i++) { + dims[i] = inputs[0]->dims[i]; + } + outputs[0] = model.create_parallel_tensor_legion_ordering( + numdim, dims, _input->data_type, this, 0 /*owner_idx*/); + + // Second output: argmax results (collapse vocab dimension to 1, keep batch dimensions) + for (int i = 1; i < numdim; i++) { + dims[i-1] = inputs[0]->dims[i]; // Shift batch dimensions down + } + int argmax_numdim = numdim - 1; // Remove vocab dimension + dims[argmax_numdim - 1].size = inputs[0]->dims[0].degree; + dims[argmax_numdim - 1].degree = inputs[0]->dims[0].degree; + dims[argmax_numdim - 1].parallel_idx = inputs[0]->dims[0].parallel_idx; + outputs[1] = model.create_parallel_tensor_legion_ordering( + argmax_numdim, dims, DT_INT32, this, 1 /*owner_idx*/); + + std::string const &input_label = remove_uid(name) + std::string(" input tensor"); + _input->print(input_label); + std::string const &softmax_label = remove_uid(name) + std::string(" softmax output tensor"); + outputs[0]->print(softmax_label); + std::string const &argmax_label = remove_uid(name) + std::string(" argmax output tensor"); + outputs[1]->print(argmax_label); +} + +Decoding::Decoding(FFModel &model, + DecodingParams const ¶ms, + const ParallelTensor input, + char const *name) + : Decoding(model, params.layer_guid, input, params.beam_search, params.name) {} + +struct DecodingInitMeta { + Decoding *decoding; +}; + +void Decoding::init_inference(FFModel const &ff, + std::vector const &batch_inputs, + std::vector const &batch_outputs, + MachineView const *mv) { + assert(check_output_input_weight_same_parallel_is()); + parallel_is = batch_outputs[0]->parallel_is; + ArgumentMap argmap; + Context ctx = ff.config.lg_ctx; + Runtime *runtime = ff.config.lg_hlr; + MachineView const *view = mv ? mv : &batch_outputs[0]->machine_view; + size_t machine_view_hash = view->hash(); + set_argumentmap_for_init_inference(ff, argmap, batch_outputs[0]); + + DecodingInitMeta meta; + meta.decoding = this; + + IndexLauncher launcher(DECODING_INIT_TASK_ID, + parallel_is, + TaskArgument(&meta, sizeof(DecodingInitMeta)), + argmap, + Predicate::TRUE_PRED, + false /*must*/, + 0 /*mapper_id*/, + machine_view_hash); + launcher.add_region_requirement(RegionRequirement(batch_inputs[0]->part, + 0 /*projection id*/, + READ_ONLY, + EXCLUSIVE, + batch_inputs[0]->region)); + launcher.add_field(0, FID_DATA); + launcher.add_region_requirement(RegionRequirement(batch_outputs[0]->part, + 0 /*projection id*/, + WRITE_DISCARD, + EXCLUSIVE, + batch_outputs[0]->region)); + launcher.add_field(1, FID_DATA); + launcher.add_region_requirement(RegionRequirement(batch_outputs[1]->part, + 0 /*projection id*/, + WRITE_DISCARD, + EXCLUSIVE, + batch_outputs[1]->region)); + launcher.add_field(2, FID_DATA); + FutureMap fm = runtime->execute_index_space(ctx, launcher); + fm.wait_all_results(); + set_opmeta_from_futuremap_inference(ff, fm, batch_outputs[0]); +} + +void Decoding::init(FFModel const &ff) { + assert(check_output_input_weight_same_parallel_is()); + parallel_is = outputs[0]->parallel_is; + ArgumentMap argmap; + Context ctx = ff.config.lg_ctx; + Runtime *runtime = ff.config.lg_hlr; + set_argumentmap_for_init(ff, argmap); + IndexLauncher launcher(DECODING_INIT_TASK_ID, + parallel_is, + TaskArgument(this, sizeof(Decoding)), + argmap, + Predicate::TRUE_PRED, + false /*must*/, + 0 /*mapper_id*/, + outputs[0]->machine_view.hash()); + launcher.add_region_requirement(RegionRequirement(inputs[0]->part, + 0 /*projection id*/, + READ_ONLY, + EXCLUSIVE, + inputs[0]->region)); + launcher.add_field(0, FID_DATA); + launcher.add_region_requirement(RegionRequirement(outputs[0]->part, + 0 /*projection id*/, + WRITE_DISCARD, + EXCLUSIVE, + outputs[0]->region)); + launcher.add_field(1, FID_DATA); + FutureMap fm = runtime->execute_index_space(ctx, launcher); + fm.wait_all_results(); + set_opmeta_from_futuremap(ff, fm); +} + +/* + regions[0]: input + regions[1]: softmax output + regions[2]: argmax output + */ +OpMeta *Decoding::init_task(Task const *task, + std::vector const ®ions, + Context ctx, + Runtime *runtime) { + assert(regions.size() == 3); + assert(task->regions.size() == regions.size()); + DecodingInitMeta const *meta = (DecodingInitMeta *)task->args; + Decoding const *decoding = meta->decoding; + + FFHandler handle = *((FFHandler const *)task->local_args); + Memory gpu_mem = get_proc_mem(Machine::get_machine(), task->target_proc); + MemoryAllocator gpu_mem_allocator(gpu_mem); + + Domain input_domain = runtime->get_index_space_domain( + ctx, task->regions[0].region.get_index_space()); + Domain softmax_output_domain = runtime->get_index_space_domain( + ctx, task->regions[1].region.get_index_space()); + Domain argmax_output_domain = runtime->get_index_space_domain( + ctx, task->regions[2].region.get_index_space()); + // Note: softmax_output_domain should match input_domain, argmax_output_domain has one fewer dimension + assert(input_domain == softmax_output_domain); + int ndims = input_domain.get_dim(); + int output_ndims = ndims - 1; // Argmax output has one fewer dimension + Domain domain; + for (int i = 0; i < ndims - 1; i++) { + assert(!decoding->outputs[0]->dims[i].is_replica_dim); + } + // Only the outter-most dim can be a replica_dim + if (decoding->outputs[0]->dims[ndims - 1].is_replica_dim) { + int replica_degree = decoding->outputs[0]->dims[ndims - 1].size; + domain.dim = ndims - 1; + for (int i = 0; i < ndims - 1; i++) { + domain.rect_data[i] = input_domain.rect_data[i]; + domain.rect_data[i + ndims - 1] = input_domain.rect_data[i + ndims]; + } + domain.rect_data[2 * ndims - 3] = + (domain.rect_data[2 * ndims - 3] + 1) * replica_degree - 1; + assert(domain.get_volume() == input_domain.get_volume()); + } else { + domain = input_domain; + } + + DecodingMeta *m = + new DecodingMeta(handle, decoding, domain, gpu_mem_allocator); + std::strcpy(m->op_name, decoding->name); + m->layer_guid = decoding->layer_guid; + m->beam_search = decoding->beam_search; + return m; +} + +void Decoding::forward(FFModel const &ff) { + // Decoding does not support forward + assert(false); +} + +FutureMap Decoding::inference(FFModel const &ff, + BatchConfigFuture const &bc, + std::vector const &batch_inputs, + std::vector const &batch_outputs, + MachineView const *mv) { + ArgumentMap argmap; + Context ctx = ff.config.lg_ctx; + Runtime *runtime = ff.config.lg_hlr; + parallel_is = batch_outputs[0]->parallel_is; + MachineView const *view = mv ? mv : &batch_outputs[0]->machine_view; + set_argumentmap_for_inference(ff, argmap, batch_outputs[0]); + size_t machine_view_hash = view->hash(); + + assert(ff.config.computationMode == COMP_MODE_INFERENCE); + + if (beam_search) { + IndexLauncher launcher(DECODING_BEAM_INF_TASK_ID, + parallel_is, + TaskArgument(nullptr, 0), + argmap, + Predicate::TRUE_PRED, + false /*must*/, + 0 /*mapper_id*/, + machine_view_hash); + launcher.add_future(bc); + launcher.add_region_requirement(RegionRequirement(batch_inputs[0]->part, + 0 /*projection id*/, + READ_ONLY, + EXCLUSIVE, + batch_inputs[0]->region)); + launcher.add_field(0, FID_DATA); + launcher.add_region_requirement(RegionRequirement(batch_outputs[0]->part, + 0 /*projection id*/, + WRITE_ONLY, + EXCLUSIVE, + batch_outputs[0]->region)); + launcher.add_field(1, FID_DATA); + launcher.add_region_requirement(RegionRequirement(batch_outputs[1]->part, + 0 /*projection id*/, + WRITE_ONLY, + EXCLUSIVE, + batch_outputs[1]->region)); + launcher.add_field(2, FID_DATA); + return runtime->execute_index_space(ctx, launcher); + } else { + IndexLauncher launcher(DECODING_NORM_INF_TASK_ID, + parallel_is, + TaskArgument(nullptr, 0), + argmap, + Predicate::TRUE_PRED, + false /*must*/, + 0 /*mapper_id*/, + machine_view_hash); + launcher.add_future(bc); + launcher.add_region_requirement(RegionRequirement(batch_inputs[0]->part, + 0 /*projection id*/, + READ_ONLY, + EXCLUSIVE, + batch_inputs[0]->region)); + launcher.add_field(0, FID_DATA); + launcher.add_region_requirement(RegionRequirement(batch_outputs[0]->part, + 0 /*projection id*/, + WRITE_ONLY, + EXCLUSIVE, + batch_outputs[0]->region)); + launcher.add_field(1, FID_DATA); + launcher.add_region_requirement(RegionRequirement(batch_outputs[1]->part, + 0 /*projection id*/, + WRITE_ONLY, + EXCLUSIVE, + batch_outputs[1]->region)); + launcher.add_field(2, FID_DATA); + return runtime->execute_index_space(ctx, launcher); + } +} + +BeamInferenceResult + Decoding::inference_task_beam(Task const *task, + std::vector const ®ions, + Context ctx, + Runtime *runtime) { + assert(regions.size() == 3); + assert(task->regions.size() == 3); + BatchConfig const *bc = BatchConfig::from_future(task->futures[0]); + if (bc->num_tokens == 0) { + // Directly return for empty batch config + BeamInferenceResult ir; + return ir; + } + DecodingMeta *m = *((DecodingMeta **)task->local_args); + + GenericTensorAccessorR input = helperGetGenericTensorAccessorRO( + m->input_type[0], regions[0], task->regions[0], FID_DATA, ctx, runtime); + GenericTensorAccessorW softmax_output = helperGetGenericTensorAccessorWO( + m->output_type[0], regions[1], task->regions[1], FID_DATA, ctx, runtime); + GenericTensorAccessorW argmax_output = helperGetGenericTensorAccessorWO( + m->output_type[1], regions[2], task->regions[2], FID_DATA, ctx, runtime); + int batch_size = bc->num_active_tokens(); + float loss = 0.0f; + + inference_kernel_wrapper(m, bc, input, softmax_output, argmax_output); + + BeamInferenceResult ir; + // Copy argmax results from output region + copy_tensor_dev_to_host( + argmax_output.get_int32_ptr(), ir.token_ids, batch_size); + copy_tensor_dev_to_host(m->probs, ir.probs, batch_size); + // Copy parent results from temporary buffer in DecodingMeta + copy_tensor_dev_to_host( + m->parent_output_buffer, ir.parent_id, batch_size); + + if (m->inference_debugging) { + assert(task->index_point.get_dim() == 1); + int shard_id = task->index_point.point_data[0]; + // Save inference tensors to file (implementation needed) + // Decoding::save_inference_tensors_to_file( + // m, shard_id, bc, {}, {}, {input, argmax_output}); + } + + return ir; +} + +InferenceResult + Decoding::inference_task_norm(Task const *task, + std::vector const ®ions, + Context ctx, + Runtime *runtime) { + assert(regions.size() == 3); + assert(task->regions.size() == 3); + DecodingMeta *m = *((DecodingMeta **)task->local_args); + BatchConfig const *bc = BatchConfig::from_future(task->futures[0]); + if (bc->num_tokens == 0) { + // Directly return for empty batch config + InferenceResult ir; + return ir; + } + + GenericTensorAccessorR input = helperGetGenericTensorAccessorRO( + m->input_type[0], regions[0], task->regions[0], FID_DATA, ctx, runtime); + GenericTensorAccessorW softmax_output = helperGetGenericTensorAccessorWO( + m->output_type[0], regions[1], task->regions[1], FID_DATA, ctx, runtime); + GenericTensorAccessorW argmax_output = helperGetGenericTensorAccessorWO( + m->output_type[1], regions[2], task->regions[2], FID_DATA, ctx, runtime); + int batch_size = bc->num_active_tokens(); + float loss = 0.0f; + + inference_kernel_wrapper(m, bc, input, softmax_output, argmax_output); + + if (task->index_point.point_data[0] == 0) { + int in_dim0 = input.domain.hi()[0] - input.domain.lo()[0] + 1; + int in_dim1 = input.domain.hi()[1] - input.domain.lo()[1] + 1; + int softmax_out_dim0 = softmax_output.domain.hi()[0] - softmax_output.domain.lo()[0] + 1; + int softmax_out_dim1 = softmax_output.domain.hi()[1] - softmax_output.domain.lo()[1] + 1; + int argmax_out_dim0 = argmax_output.domain.hi()[0] - argmax_output.domain.lo()[0] + 1; + std::string op_name_without_uid = remove_uid(m->op_name); + printf("Decoding(%s): in=[%i, bz=%i/%i] -> softmax_out=[%i,bz=%i/%i], argmax_out=[bz=%i]\n", + op_name_without_uid.c_str(), + in_dim0, bc->num_tokens, in_dim1, + softmax_out_dim0, bc->num_tokens, softmax_out_dim1, + argmax_out_dim0); + } + + InferenceResult ir; + ir.finetuning_loss = loss; + + if (m->inference_debugging) { + assert(task->index_point.get_dim() == 1); + int shard_id = task->index_point.point_data[0]; + // Save inference tensors to file (implementation needed) + Decoding::save_inference_tensors_to_file( + m, shard_id, bc, {input}, {}, {softmax_output, argmax_output}); + } else { + m->decoding_step++; + } + + // Copy argmax results from output region + copy_tensor_dev_to_host( + argmax_output.get_int32_ptr(), ir.token_ids, batch_size); + + return ir; +} + +void Decoding::backward(FFModel const &ff) { + // Decoding does not support backward + assert(false); +} + +FutureMap Decoding::peft_bwd(FFModel const &ff, + BatchConfigFuture const &bc, + std::vector const &batch_inputs, + std::vector const &batch_outputs, + MachineView const *mv) { + ArgumentMap argmap; + Context ctx = ff.config.lg_ctx; + Runtime *runtime = ff.config.lg_hlr; + parallel_is = batch_outputs[0]->parallel_is; + MachineView const *view = mv ? mv : &batch_outputs[0]->machine_view; + set_argumentmap_for_inference(ff, argmap, batch_outputs[0]); + size_t machine_view_hash = view->hash(); + IndexLauncher launcher(DECODING_PEFT_BWD_TASK_ID, + parallel_is, + TaskArgument(nullptr, 0), + argmap, + Predicate::TRUE_PRED, + false /*must*/, + 0 /*mapper_id*/, + machine_view_hash); + launcher.add_future(bc); + launcher.add_region_requirement( + RegionRequirement(batch_inputs[0]->part_grad, + 0 /*projection id*/, + reset_input_grads[0] ? WRITE_ONLY : READ_WRITE, + EXCLUSIVE, + batch_inputs[0]->region_grad)); + launcher.add_field(0, FID_DATA); + return runtime->execute_index_space(ctx, launcher); +} + +bool Decoding::peft_bwd_task(Task const *task, + std::vector const ®ions, + Context ctx, + Runtime *runtime) { + assert(task->regions.size() == regions.size()); + assert(regions.size() == 1); + assert(task->regions.size() == 1); + BatchConfig const *bc = BatchConfig::from_future(task->futures[0]); + DecodingMeta *m = *((DecodingMeta **)task->local_args); + if (!bc->peft_bwd_applies_to_this_layer(m->layer_guid.transformer_layer_id)) { + return false; + } + Domain in_domain = runtime->get_index_space_domain( + ctx, task->regions[0].region.get_index_space()); + + GenericTensorAccessorW input_grad = helperGetGenericTensorAccessorRW( + m->input_type[0], regions[0], task->regions[0], FID_DATA, ctx, runtime); + + peft_bwd_kernel_wrapper(m, bc, task->index_point.point_data[0], input_grad); + if (m->inference_debugging) { + assert(task->index_point.get_dim() == 1); + int shard_id = task->index_point.point_data[0]; + // Save inference tensors to file (implementation needed) + Decoding::save_inference_tensors_to_file( + m, shard_id, bc, {input_grad}, {}, {}, false); + } + return true; +} + +bool Decoding::measure_operator_cost(Simulator *sim, + MachineView const &mv, + CostMetrics &cost_metrics) const { + return false; +} + +Op *Decoding::materialize(FFModel &ff, + ParallelTensor inputs[], + int num_inputs) const { + DecodingParams params = get_params(); + return new Decoding(ff, params, inputs[0], this->name); +} + +}; // namespace FlexFlow + +namespace std { +size_t hash::operator()( + FlexFlow::DecodingParams const ¶ms) const { + size_t key = 0; + hash_combine(key, params.layer_guid.id); + hash_combine(key, params.beam_search); + return key; +} +}; // namespace std diff --git a/src/ops/decoding.cu b/src/ops/decoding.cu new file mode 100644 index 000000000..5a9779b90 --- /dev/null +++ b/src/ops/decoding.cu @@ -0,0 +1,792 @@ +/* Copyright 2023 CMU, Facebook, LANL, MIT, NVIDIA, and Stanford (alphabetical) + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "flexflow/ffconst_utils.h" +#include "flexflow/ops/decoding.h" +#include "flexflow/utils/cuda_helper.h" + +namespace FlexFlow { + +template +__global__ void softmax_argmax_kernel_sharded( + const half* __restrict__ input, + half* __restrict__ output, + float* __restrict__ max_buffer, // [seq_len] for allreduce + int* __restrict__ max_idx_buffer, // [seq_len] for argmax indices + float* __restrict__ sum_buffer, // [seq_len] for allreduce + const int vocab_size_per_shard, + const int vocab_offset, // Starting vocab index for this shard + const int seq_len) { + + const int seq_pos = blockIdx.x; + if (seq_pos >= seq_len) return; + + const int tid = threadIdx.x; + const int lane_id = tid % 32; + const int warp_id = tid / 32; + const int num_warps = BLOCK_SIZE / 32; + + extern __shared__ char shared_mem_bytes[]; + float* warp_max = (float*)shared_mem_bytes; + int* warp_max_idx = (int*)(warp_max + num_warps); + float* warp_sum = (float*)(warp_max_idx + num_warps); + + // Phase 1: Find local max and its index + float thread_max = -INFINITY; + int thread_max_idx = 0; + + for (int i = tid; i < vocab_size_per_shard; i += BLOCK_SIZE) { + float val = __half2float(input[seq_pos * vocab_size_per_shard + i]); + if (val > thread_max) { + thread_max = val; + thread_max_idx = vocab_offset + i; // Global vocab index + } + } + + // Warp reduction for max with index + #pragma unroll + for (int offset = 16; offset > 0; offset /= 2) { + float other_max = __shfl_down_sync(0xffffffff, thread_max, offset); + int other_idx = __shfl_down_sync(0xffffffff, thread_max_idx, offset); + + if (other_max > thread_max || (other_max == thread_max && other_idx < thread_max_idx)) { + thread_max = other_max; + thread_max_idx = other_idx; + } + } + + if (lane_id == 0) { + warp_max[warp_id] = thread_max; + warp_max_idx[warp_id] = thread_max_idx; + } + __syncthreads(); + + // Final reduction across warps + if (tid < num_warps) { + thread_max = warp_max[tid]; + thread_max_idx = warp_max_idx[tid]; + } else { + thread_max = -INFINITY; + thread_max_idx = INT_MAX; + } + + if (warp_id == 0) { + #pragma unroll + for (int offset = 16; offset > 0; offset /= 2) { + float other_max = __shfl_down_sync(0xffffffff, thread_max, offset); + int other_idx = __shfl_down_sync(0xffffffff, thread_max_idx, offset); + + if (other_max > thread_max || (other_max == thread_max && other_idx < thread_max_idx)) { + thread_max = other_max; + thread_max_idx = other_idx; + } + } + + if (lane_id == 0) { + max_buffer[seq_pos] = thread_max; + max_idx_buffer[seq_pos] = thread_max_idx; + } + } +} + +// Custom reduction operation for max with index +__device__ void reduce_max_with_idx(float* max_val, int* max_idx, float other_val, int other_idx) { + if (other_val > *max_val || (other_val == *max_val && other_idx < *max_idx)) { + *max_val = other_val; + *max_idx = other_idx; + } +} + +template +__global__ void softmax_compute_kernel_sharded( + const half* __restrict__ input, + half* __restrict__ output, + const float* __restrict__ max_buffer, // Global max from allreduce + float* __restrict__ sum_buffer, + const int vocab_size_per_shard, + const int seq_len) { + + const int seq_pos = blockIdx.x; + if (seq_pos >= seq_len) return; + + const int tid = threadIdx.x; + const int lane_id = tid % 32; + const int warp_id = tid / 32; + const int num_warps = BLOCK_SIZE / 32; + + extern __shared__ float warp_sum[]; + + float max_val = max_buffer[seq_pos]; + float thread_sum = 0.0f; + + // Compute exp(x - max) and local sum + for (int i = tid; i < vocab_size_per_shard; i += BLOCK_SIZE) { + float val = __half2float(input[seq_pos * vocab_size_per_shard + i]); + float exp_val = expf(val - max_val); + thread_sum += exp_val; + output[seq_pos * vocab_size_per_shard + i] = __float2half(exp_val); + } + + // Warp reduction for sum + #pragma unroll + for (int offset = 16; offset > 0; offset /= 2) { + thread_sum += __shfl_down_sync(0xffffffff, thread_sum, offset); + } + + if (lane_id == 0) warp_sum[warp_id] = thread_sum; + __syncthreads(); + + // Final reduction + if (tid < num_warps) thread_sum = warp_sum[tid]; + else thread_sum = 0.0f; + + if (warp_id == 0) { + #pragma unroll + for (int offset = 16; offset > 0; offset /= 2) { + thread_sum += __shfl_down_sync(0xffffffff, thread_sum, offset); + } + if (lane_id == 0) sum_buffer[seq_pos] = thread_sum; + } +} + +template +__global__ void softmax_normalize_kernel_sharded( + half* __restrict__ output, + const float* __restrict__ sum_buffer, // Global sum from allreduce + const int vocab_size_per_shard, + const int seq_len) { + + const int seq_pos = blockIdx.x; + if (seq_pos >= seq_len) return; + + const int tid = threadIdx.x; + float inv_sum = 1.0f / sum_buffer[seq_pos]; + + // Normalize + for (int i = tid; i < vocab_size_per_shard; i += BLOCK_SIZE) { + float exp_val = __half2float(output[seq_pos * vocab_size_per_shard + i]); + output[seq_pos * vocab_size_per_shard + i] = __float2half(exp_val * inv_sum); + } +} + +// Kernel to handle custom reduction for max with argmax +__global__ void reduce_max_with_argmax_kernel( + float* max_buffer, + int* max_idx_buffer, + int seq_len, + int n_shards) { + + int seq_pos = blockIdx.x * blockDim.x + threadIdx.x; + if (seq_pos >= seq_len) return; + + // Each thread handles one sequence position + // The buffers contain [shard0_val, shard1_val, ...] for each sequence + float global_max = max_buffer[seq_pos * n_shards]; + int global_idx = max_idx_buffer[seq_pos * n_shards]; + + for (int shard = 1; shard < n_shards; shard++) { + float shard_max = max_buffer[seq_pos * n_shards + shard]; + int shard_idx = max_idx_buffer[seq_pos * n_shards + shard]; + + if (shard_max > global_max || (shard_max == global_max && shard_idx < global_idx)) { + global_max = shard_max; + global_idx = shard_idx; + } + } + + // Write back the global max and index + max_buffer[seq_pos] = global_max; + max_idx_buffer[seq_pos] = global_idx; +} + +struct SoftmaxShardedContext { + float* max_buffer; + int* max_idx_buffer; + float* sum_buffer; + float* all_max_buffer; + int* all_idx_buffer; + int seq_len; + int n_shards; +}; + +SoftmaxShardedContext* create_softmax_context(int seq_len, ncclComm_t nccl_comm) { + SoftmaxShardedContext* ctx = new SoftmaxShardedContext; + ctx->seq_len = seq_len; + ncclCommCount(nccl_comm, &ctx->n_shards); + + cudaMalloc(&ctx->max_buffer, seq_len * sizeof(float)); + cudaMalloc(&ctx->max_idx_buffer, seq_len * sizeof(int)); + cudaMalloc(&ctx->sum_buffer, seq_len * sizeof(float)); + cudaMalloc(&ctx->all_max_buffer, seq_len * ctx->n_shards * sizeof(float)); + cudaMalloc(&ctx->all_idx_buffer, seq_len * ctx->n_shards * sizeof(int)); + + return ctx; +} + +void softmax_argmax_sharded_with_context( + const half* input, + half* output, + int* argmax_indices, + int vocab_size_per_shard, + int vocab_offset, + SoftmaxShardedContext* ctx, + cudaStream_t stream, + ncclComm_t nccl_comm) { + + const int block_size = 256; + const int grid_size = ctx->seq_len; + size_t shared_mem_size = (block_size / 32) * sizeof(float) * 2 + (block_size / 32) * sizeof(int); + + // Step 1: Find local max and argmax + softmax_argmax_kernel_sharded<<>>( + input, output, ctx->max_buffer, ctx->max_idx_buffer, ctx->sum_buffer, + vocab_size_per_shard, vocab_offset, ctx->seq_len + ); + + // Steps 2-5: Same as before but using pre-allocated buffers from context + ncclAllReduce(ctx->max_buffer, ctx->max_buffer, ctx->seq_len, ncclFloat32, ncclMax, nccl_comm, stream); + + ncclAllGather(ctx->max_buffer, ctx->all_max_buffer, ctx->seq_len, ncclFloat32, nccl_comm, stream); + ncclAllGather(ctx->max_idx_buffer, ctx->all_idx_buffer, ctx->seq_len, ncclInt32, nccl_comm, stream); + + int reduce_blocks = (ctx->seq_len + 255) / 256; + reduce_max_with_argmax_kernel<<>>( + ctx->all_max_buffer, ctx->all_idx_buffer, ctx->seq_len, ctx->n_shards + ); + + cudaMemcpyAsync(argmax_indices, ctx->all_idx_buffer, ctx->seq_len * sizeof(int), + cudaMemcpyDeviceToDevice, stream); + + softmax_compute_kernel_sharded<<>>( + input, output, ctx->max_buffer, ctx->sum_buffer, vocab_size_per_shard, ctx->seq_len + ); + + ncclAllReduce(ctx->sum_buffer, ctx->sum_buffer, ctx->seq_len, ncclFloat32, ncclSum, nccl_comm, stream); + + softmax_normalize_kernel_sharded<<>>( + output, ctx->sum_buffer, vocab_size_per_shard, ctx->seq_len + ); +} + +void destroy_softmax_context(SoftmaxShardedContext* ctx) { + cudaFree(ctx->max_buffer); + cudaFree(ctx->max_idx_buffer); + cudaFree(ctx->sum_buffer); + cudaFree(ctx->all_max_buffer); + cudaFree(ctx->all_idx_buffer); + delete ctx; +} + +// Optimized version using warp primitives +template +__global__ void softmax_argmax_kernel( + const DT* __restrict__ input, + DT* __restrict__ output, + int* __restrict__ argmax_indices, + const int vocab_size, + const int seq_len) { + + const int seq_pos = blockIdx.x; + if (seq_pos >= seq_len) return; + + const int tid = threadIdx.x; + const int lane_id = tid % 32; + const int warp_id = tid / 32; + const int num_warps = BLOCK_SIZE / 32; + + // Shared memory + extern __shared__ char shared_mem_bytes[]; + float* warp_max = (float*)shared_mem_bytes; + int* warp_max_idx = (int*)(warp_max + num_warps); + float* warp_sum = (float*)(warp_max_idx + num_warps); + + // Phase 1: Find max + float thread_max = -INFINITY; + int thread_max_idx = 0; + + for (int vocab_idx = tid; vocab_idx < vocab_size; vocab_idx += BLOCK_SIZE) { + float val; + if constexpr (std::is_same_v) { + val = __half2float(input[seq_pos * vocab_size + vocab_idx]); + } else { + val = static_cast(input[seq_pos * vocab_size + vocab_idx]); + } + if (val > thread_max) { + thread_max = val; + thread_max_idx = vocab_idx; + } + } + + // Warp-level reduction for max + #pragma unroll + for (int offset = 16; offset > 0; offset /= 2) { + float other_max = __shfl_down_sync(0xffffffff, thread_max, offset); + int other_idx = __shfl_down_sync(0xffffffff, thread_max_idx, offset); + + if (other_max > thread_max || (other_max == thread_max && other_idx < thread_max_idx)) { + thread_max = other_max; + thread_max_idx = other_idx; + } + } + + // Store warp results + if (lane_id == 0) { + warp_max[warp_id] = thread_max; + warp_max_idx[warp_id] = thread_max_idx; + } + __syncthreads(); + + // Final reduction across warps + if (tid < num_warps) { + thread_max = warp_max[tid]; + thread_max_idx = warp_max_idx[tid]; + } else { + thread_max = -INFINITY; + thread_max_idx = INT_MAX; + } + + if (warp_id == 0) { + #pragma unroll + for (int offset = 16; offset > 0; offset /= 2) { + float other_max = __shfl_down_sync(0xffffffff, thread_max, offset); + int other_idx = __shfl_down_sync(0xffffffff, thread_max_idx, offset); + + if (other_max > thread_max || (other_max == thread_max && other_idx < thread_max_idx)) { + thread_max = other_max; + thread_max_idx = other_idx; + } + } + + if (lane_id == 0) { + warp_max[0] = thread_max; + warp_max_idx[0] = thread_max_idx; + } + } + __syncthreads(); + + float max_val = warp_max[0]; + int max_idx = warp_max_idx[0]; + + // Phase 2: Compute exp and sum + float thread_sum = 0.0f; + + for (int vocab_idx = tid; vocab_idx < vocab_size; vocab_idx += BLOCK_SIZE) { + float val; + if constexpr (std::is_same_v) { + val = __half2float(input[seq_pos * vocab_size + vocab_idx]); + } else { + val = static_cast(input[seq_pos * vocab_size + vocab_idx]); + } + float exp_val = expf(val - max_val); + thread_sum += exp_val; + if constexpr (std::is_same_v) { + output[seq_pos * vocab_size + vocab_idx] = __float2half(exp_val); + } else { + output[seq_pos * vocab_size + vocab_idx] = static_cast
(exp_val); + } + } + + // Warp-level reduction for sum + #pragma unroll + for (int offset = 16; offset > 0; offset /= 2) { + thread_sum += __shfl_down_sync(0xffffffff, thread_sum, offset); + } + + if (lane_id == 0) { + warp_sum[warp_id] = thread_sum; + } + __syncthreads(); + + // Final reduction + if (tid < num_warps) { + thread_sum = warp_sum[tid]; + } else { + thread_sum = 0.0f; + } + + if (warp_id == 0) { + #pragma unroll + for (int offset = 16; offset > 0; offset /= 2) { + thread_sum += __shfl_down_sync(0xffffffff, thread_sum, offset); + } + + if (lane_id == 0) { + warp_sum[0] = thread_sum; + } + } + __syncthreads(); + + float total_sum = warp_sum[0]; + float inv_sum = 1.0f / total_sum; + + // Phase 3: Normalize + for (int vocab_idx = tid; vocab_idx < vocab_size; vocab_idx += BLOCK_SIZE) { + float exp_val; + if constexpr (std::is_same_v) { + exp_val = __half2float(output[seq_pos * vocab_size + vocab_idx]); + output[seq_pos * vocab_size + vocab_idx] = __float2half(exp_val * inv_sum); + } else { + exp_val = static_cast(output[seq_pos * vocab_size + vocab_idx]); + output[seq_pos * vocab_size + vocab_idx] = static_cast
(exp_val * inv_sum); + } + } + +// Store argmax + if (tid == 0) { + argmax_indices[seq_pos] = max_idx; + } +} + +// Wrapper function template +template +void softmax_argmax( + const DT* input, + DT* output, + int* argmax_indices, + int vocab_size, + int seq_len, + cudaStream_t stream) { + + const int block_size = 256; + const int grid_size = seq_len; // One block per sequence position + + + // Use optimized kernel for larger vocabularies + size_t shared_mem_size_opt = (block_size / 32) * sizeof(float) * 2 + (block_size / 32) * sizeof(int); + softmax_argmax_kernel<<>>( + input, output, argmax_indices, vocab_size, seq_len + ); +} + +// Placeholder kernel implementations +template +void Decoding::inference_kernel(DecodingMeta const *m, + BatchConfig const *bc, + DT const *input_ptr, + DT *softmax_output_ptr, + int *argmax_output_ptr, + int num_classes, + int vocab_offset, + float *loss, + cudaStream_t stream) { + + // Old non-sharded version (commented out) + // softmax_argmax(input_ptr, softmax_output_ptr, argmax_output_ptr, num_classes, bc->num_active_tokens(), stream); + + // New sharded version - only supports half precision for now + if constexpr (std::is_same_v) { + // Use sharded softmax with NCCL communication + int vocab_size_per_shard = num_classes; // This is already the local shard size + + softmax_argmax_sharded_with_context( + input_ptr, + softmax_output_ptr, + argmax_output_ptr, + vocab_size_per_shard, + vocab_offset, + m->softmax_context, + stream, + m->handle.ncclComm); + } else { + // Fall back to non-sharded version for non-half types + softmax_argmax(input_ptr, softmax_output_ptr, argmax_output_ptr, num_classes, bc->num_active_tokens(), stream); + } +} + +void store_peft_token_ids(DecodingMeta *m, BatchConfig const *bc) { + assert(peft_finetuning_enabled(m->peft_support_mode)); + + int num_ft_tokens = bc->num_finetuning_fwd_tokens(); + int i = bc->finetuning_request_index(); + int tokens_previous_requests = + bc->requestsInfo[i].first_token_offset_in_batch; + int prev_steps_tokens = bc->requestsInfo[i].first_token_depth_in_request; + assert(bc->requestsInfo[i].num_tokens_in_batch == num_ft_tokens); + + // shift labels by 1 position to the left (ignore first token label) + for (int j = 0; j < num_ft_tokens - 1; j++) { + m->peft_token_ids[prev_steps_tokens + j] = + bc->tokensInfo[tokens_previous_requests + j + 1].token_id; + } +} + +template +void store_peft_activations(DecodingMeta *m, + BatchConfig const *bc, + int num_classes, + DT *softmax_output_ptr, + cudaStream_t stream) { + assert(peft_finetuning_enabled(m->peft_support_mode)); + assert(m->output_grad_ptr != nullptr); + + int num_ft_tokens = bc->num_finetuning_fwd_tokens(); + int i = bc->finetuning_request_index(); + int tokens_previous_requests = + bc->requestsInfo[i].first_token_offset_in_batch; + int prev_steps_tokens = bc->requestsInfo[i].first_token_depth_in_request; + assert(bc->requestsInfo[i].num_tokens_in_batch == num_ft_tokens); + + size_t batch_offset = num_classes * tokens_previous_requests; + size_t req_offset = num_classes * prev_steps_tokens; + size_t data_size = num_classes * num_ft_tokens * sizeof(DT); + assert(m->allocated_peft_buffer_size >= data_size); + checkCUDA(cudaMemcpyAsync(static_cast
(m->output_grad_ptr) + req_offset, + softmax_output_ptr + batch_offset, + data_size, + cudaMemcpyDeviceToDevice, + stream)); +} + +/*static*/ +void Decoding::inference_kernel_wrapper(DecodingMeta *m, + BatchConfig const *bc, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &softmax_output, + GenericTensorAccessorW const &argmax_output) { + cudaStream_t stream; + checkCUDA(get_legion_stream(&stream)); + cudaEvent_t t_start, t_end; + if (m->profiling) { + cudaEventCreate(&t_start); + cudaEventCreate(&t_end); + cudaEventRecord(t_start, stream); + } + + if (bc->num_active_tokens() <= 0) { + return; + } + + int num_classes = input.domain.hi()[0] - input.domain.lo()[0] + 1; + int vocab_offset = input.domain.lo()[0]; // Starting vocab index for this shard + float loss = 0.0f; + + if (input.data_type == DT_HALF) { + Decoding::inference_kernel(m, + bc, + input.get_half_ptr(), + softmax_output.get_half_ptr(), + argmax_output.get_int32_ptr(), + num_classes, + vocab_offset, + &loss, + stream); + } else if (input.data_type == DT_FLOAT) { + Decoding::inference_kernel(m, + bc, + input.get_float_ptr(), + softmax_output.get_float_ptr(), + argmax_output.get_int32_ptr(), + num_classes, + vocab_offset, + &loss, + stream); + } else { + assert(false && "Unsupported data type"); + } + + if (bc->num_finetuning_fwd_requests() > 0) { + store_peft_token_ids(m, bc); + // Store softmax activations for PEFT backward pass + if (input.data_type == DT_HALF) { + store_peft_activations(m, bc, num_classes, softmax_output.get_half_ptr(), stream); + } else if (input.data_type == DT_FLOAT) { + store_peft_activations(m, bc, num_classes, softmax_output.get_float_ptr(), stream); + } + } + + if (m->profiling) { + cudaEventRecord(t_end, stream); + checkCUDA(cudaEventSynchronize(t_end)); + float elapsed = 0; + checkCUDA(cudaEventElapsedTime(&elapsed, t_start, t_end)); + cudaEventDestroy(t_start); + cudaEventDestroy(t_end); + printf("[Decoding] forward time = %.2lfms\n", elapsed); + } +} + +template +__global__ void sparse_categorical_crossentropy_loss_peft_backward( + DT *input_grad, + DT const *output_grad, + BatchConfig::TokenId const *token_ids, + int num_tokens, + int num_classes, + int shard_id) { + CUDA_KERNEL_LOOP(i, num_tokens * num_classes) { + int class_idx = i % num_classes; + int token_idx = i / num_classes; + input_grad[i] = output_grad[i]; + if (class_idx + shard_id * num_classes == token_ids[token_idx]) { + input_grad[i] = input_grad[i] - (DT)1.0f; + } + } +} + +template +void Decoding::peft_bwd_kernel(DecodingMeta const *m, + BatchConfig const *bc, + DT *input_grad_ptr, + int num_classes, + int shard_id, + cudaStream_t stream) { + printf("peft_bwd_kernel - num_classes: %d, shard_id: %d\n", num_classes, shard_id); + assert( + bc->peft_bwd_applies_to_this_layer(m->layer_guid.transformer_layer_id)); + int i = bc->finetuning_request_index(); + + int num_bwd_tokens = bc->requestsInfo[i].num_tokens_in_batch - 1; + + DT scale_factor = 1.0 / (bc->requestsInfo[i].num_tokens_in_batch); + // ignore last token + checkCUDA(cudaMemsetAsync(input_grad_ptr + num_bwd_tokens * num_classes, + 0, + num_classes * sizeof(DT), + stream)); + checkCUDA(cudaMemcpyAsync(m->handle.workSpace, + m->peft_token_ids, + sizeof(BatchConfig::TokenId) * num_bwd_tokens, + cudaMemcpyHostToDevice, + stream)); + sparse_categorical_crossentropy_loss_peft_backward<<< + GET_BLOCKS(num_bwd_tokens * num_classes), + CUDA_NUM_THREADS, + 0, + stream>>>(input_grad_ptr, + static_cast
(m->output_grad_ptr), + static_cast(m->handle.workSpace), + num_bwd_tokens, + num_classes, + shard_id); + // scale + scale_kernel<<>>( + input_grad_ptr, num_bwd_tokens * num_classes, DT(0.0), scale_factor); +} + +/*static*/ +void Decoding::peft_bwd_kernel_wrapper(DecodingMeta *m, + BatchConfig const *bc, + int shard_id, + GenericTensorAccessorW const &input_grad) { + cudaStream_t stream; + checkCUDA(get_legion_stream(&stream)); + cudaEvent_t t_start, t_end; + if (m->profiling) { + cudaEventCreate(&t_start); + cudaEventCreate(&t_end); + cudaEventRecord(t_start, stream); + } + + int num_classes = input_grad.domain.hi()[0] - input_grad.domain.lo()[0] + 1; + if (m->input_type[0] == DT_FLOAT) { + Decoding::peft_bwd_kernel( + m, bc, input_grad.get_float_ptr(), num_classes, shard_id, stream); + } else if (m->input_type[0] == DT_HALF) { + Decoding::peft_bwd_kernel( + m, bc, input_grad.get_half_ptr(), num_classes, shard_id, stream); + } else { + assert(false && "Unsupported data type"); + } + if (m->profiling) { + cudaEventRecord(t_end, stream); + checkCUDA(cudaEventSynchronize(t_end)); + float elapsed = 0; + checkCUDA(cudaEventElapsedTime(&elapsed, t_start, t_end)); + cudaEventDestroy(t_start); + cudaEventDestroy(t_end); + printf("[Decoding] peft_bwd time = %.2fms\n", elapsed); + } +} + +DecodingMeta::DecodingMeta(FFHandler handler, + Decoding const *decoding, + Legion::Domain const &input_domain, + MemoryAllocator &gpu_mem_allocator) + : OpMeta(handler, decoding) { + beam_search = decoding->beam_search; + + if (peft_finetuning_enabled(peft_support_mode)) { + allocated_peft_buffer_size = + input_domain.get_volume() * data_type_size(decoding->data_type); + gpu_mem_allocator.create_legion_instance( + reserveInst, allocated_peft_buffer_size, "DecodingMeta"); + output_grad_ptr = + gpu_mem_allocator.allocate_instance_untyped(allocated_peft_buffer_size); + } else { + allocated_peft_buffer_size = 0; + output_grad_ptr = nullptr; + } + + // Simple allocations for required buffers + probs = nullptr; // Not needed for basic decoding + d_loss = nullptr; // Not needed for basic decoding + parent_output_buffer = nullptr; // Only needed for beam search if implemented + + // Create softmax context for sharded computation + // Use a reasonable max sequence length based on input domain + int max_seq_len = input_domain.get_dim() > 1 ? + (input_domain.hi()[input_domain.get_dim()-1] - input_domain.lo()[input_domain.get_dim()-1] + 1) : 1024; + softmax_context = create_softmax_context(max_seq_len, handler.ncclComm); + + std::strcpy(op_name, decoding->name); +} + +DecodingMeta::~DecodingMeta(void) { + if (reserveInst != Realm::RegionInstance::NO_INST) { + reserveInst.destroy(); + } + // Destroy softmax context + if (softmax_context != nullptr) { + destroy_softmax_context(softmax_context); + } +} + +// Explicit template instantiations +template void softmax_argmax( + const half* input, + half* output, + int* argmax_indices, + int vocab_size, + int seq_len, + cudaStream_t stream); + +template void softmax_argmax( + const float* input, + float* output, + int* argmax_indices, + int vocab_size, + int seq_len, + cudaStream_t stream); + +// Explicit template instantiations for peft_bwd_kernel +template void Decoding::peft_bwd_kernel( + DecodingMeta const *m, + BatchConfig const *bc, + half *input_grad_ptr, + int num_classes, + int shard_id, + cudaStream_t stream); + +template void Decoding::peft_bwd_kernel( + DecodingMeta const *m, + BatchConfig const *bc, + float *input_grad_ptr, + int num_classes, + int shard_id, + cudaStream_t stream); + +} // namespace FlexFlow \ No newline at end of file diff --git a/src/ops/linear.cc b/src/ops/linear.cc index c1a87c06b..9d3e07e31 100644 --- a/src/ops/linear.cc +++ b/src/ops/linear.cc @@ -117,6 +117,22 @@ Tensor FFModel::dense(const Tensor input, return li->outputs[0]; } +static std::string remove_uid(char const *op_name) { + std::string op_name_without_uid = std::string(op_name); + size_t last_underscore = op_name_without_uid.length(); + for (int i = op_name_without_uid.length() - 1; i > 0; i--) { + if (!(std::isdigit(op_name[i]) || op_name[i] == '_')) { + break; + } else if (op_name[i] == '_') { + last_underscore = i; + } + } + if (last_underscore < op_name_without_uid.length()) { + op_name_without_uid.erase(last_underscore); + } + return op_name_without_uid; +} + Op *Linear::create_operator_from_layer( FFModel &model, Layer const *layer, @@ -226,6 +242,9 @@ Linear::Linear(FFModel &model, this->in_channels = _input->dims[dimension_names.at(LinearParams::INPUT_CHANNEL)].size; + std::string const &input_label = remove_uid(name) + std::string(" input tensor"); + inputs[0]->print(input_label); + ParallelTensorShape input_shape = this->inputs[0]->get_shape(); ParallelTensorShape output_shape, kernel_shape, bias_shape; LinearParams params = this->get_params(); @@ -263,6 +282,9 @@ Linear::Linear(FFModel &model, kernel_initializer, CHOSEN_SYNC_TYPE); + std::string const &weight_label = remove_uid(name) + std::string(" weight tensor"); + weights[KERNEL_IDX]->print(weight_label); + if (use_bias) { Initializer *bias_initializer = new ZeroInitializer(); @@ -275,6 +297,9 @@ Linear::Linear(FFModel &model, bias_initializer, CHOSEN_SYNC_TYPE); add_bias_only_once = _input->dims[0].degree > 1; + + std::string const &bias_label = remove_uid(name) + std::string(" bias tensor"); + weights[BIAS_IDX]->print(bias_label); } } @@ -282,6 +307,9 @@ Linear::Linear(FFModel &model, outputs[0] = model.create_parallel_tensor_legion_ordering( output_shape.num_dims, output_shape.dims, _data_type, this); + std::string const &label = remove_uid(name) + std::string(" output tensor"); + outputs[0]->print(label); + // assert(check_output_input_weight_parallel_dims(allocate_weights)); } @@ -643,6 +671,17 @@ void Linear::inference_task(Task const *task, bias.ptr, in_dim, out_dim); + if (task->index_point.point_data[0] == 0) { + std::string op_name_without_uid = get_op_name_without_uid(m); + printf("\t%s: w=[%i,%i].T @ in=[%i, bz=%i] -> out=[%i,bz=%i]\n", + op_name_without_uid.c_str(), + in_dim, + out_dim, + in_dim, + bc->num_tokens, + out_dim, + bc->num_tokens); + } if (m->inference_debugging) { assert(task->index_point.get_dim() == 1); int shard_id = task->index_point.point_data[0]; diff --git a/src/ops/softmax.cc b/src/ops/softmax.cc index 808cbf720..d07b8aa30 100644 --- a/src/ops/softmax.cc +++ b/src/ops/softmax.cc @@ -136,6 +136,22 @@ Op *Softmax::create_operator_from_layer( layer->name); } +static std::string remove_uid(char const *op_name) { + std::string op_name_without_uid = std::string(op_name); + size_t last_underscore = op_name_without_uid.length(); + for (int i = op_name_without_uid.length() - 1; i > 0; i--) { + if (!(std::isdigit(op_name[i]) || op_name[i] == '_')) { + break; + } else if (op_name[i] == '_') { + last_underscore = i; + } + } + if (last_underscore < op_name_without_uid.length()) { + op_name_without_uid.erase(last_underscore); + } + return op_name_without_uid; +} + Softmax::Softmax(FFModel &model, LayerID const &_layer_guid, const ParallelTensor _input, @@ -159,6 +175,10 @@ Softmax::Softmax(FFModel &model, dims[i] = _input->dims[numdim - 1 - i]; } outputs[0] = model.create_parallel_tensor(numdim, dims, data_type, this); + std::string const &input_label = remove_uid(name) + std::string(" input tensor"); + _input->print(input_label); + std::string const &label = remove_uid(name) + std::string(" output tensor"); + outputs[0]->print(label); } Softmax::Softmax(FFModel &model, @@ -482,7 +502,7 @@ void Softmax::inference_task(Task const *task, m->input_type[0], regions[0], task->regions[0], FID_DATA, ctx, runtime); GenericTensorAccessorW output = helperGetGenericTensorAccessorWO( m->output_type[0], regions[1], task->regions[1], FID_DATA, ctx, runtime); - GenericTensorAccessorW output_grad; + // GenericTensorAccessorW output_grad; // if (is_last_op) { // output_grad = helperGetGenericTensorAccessorWO(m->output_type[0], // regions[2], @@ -492,6 +512,20 @@ void Softmax::inference_task(Task const *task, // runtime); // } inference_kernel_wrapper(m, bc, is_last_op, input, output); + + if (task->index_point.point_data[0] == 0) { + int in_dim0 = input.domain.hi()[0] - input.domain.lo()[0] + 1; + int in_dim1 = input.domain.hi()[1] - input.domain.lo()[1] + 1; + int out_dim0 = output.domain.hi()[0] - output.domain.lo()[0] + 1; + int out_dim1 = output.domain.hi()[1] - output.domain.lo()[1] + 1; + std::string op_name_without_uid = remove_uid(m->op_name); + printf("Softmax(%s): in=[%i, bz=%i/%i] -> out=[%i,bz=%i/%i]\n", + op_name_without_uid.c_str(), + in_dim0, bc->num_tokens, in_dim1, + out_dim0, bc->num_tokens, out_dim1); + } + + if (m->inference_debugging) { assert(task->index_point.get_dim() == 1); int shard_id = task->index_point.point_data[0]; diff --git a/src/parallel_ops/combine.cc b/src/parallel_ops/combine.cc index a47835c30..0d7d6c2ae 100644 --- a/src/parallel_ops/combine.cc +++ b/src/parallel_ops/combine.cc @@ -75,6 +75,22 @@ Combine::Combine(FFModel &model, params.combine_degree, params.name) {} +static std::string remove_uid(char const *op_name) { + std::string op_name_without_uid = std::string(op_name); + size_t last_underscore = op_name_without_uid.length(); + for (int i = op_name_without_uid.length() - 1; i > 0; i--) { + if (!(std::isdigit(op_name[i]) || op_name[i] == '_')) { + break; + } else if (op_name[i] == '_') { + last_underscore = i; + } + } + if (last_underscore < op_name_without_uid.length()) { + op_name_without_uid.erase(last_underscore); + } + return op_name_without_uid; +} + Combine::Combine(FFModel &model, const ParallelTensor _input, int _combine_legion_dim, @@ -95,6 +111,10 @@ Combine::Combine(FFModel &model, numdim, dims, _input->data_type, this); // inputs[0]->print("Combine::input"); // outputs[0]->print("Combine::output"); + std::string const &input_label = std::string("Combine input tensor"); + _input->print(input_label); + std::string const &label = std::string("Combine output tensor"); + outputs[0]->print(label); } OpMeta *Combine::init_task(Task const *task, @@ -488,6 +508,16 @@ void Combine::forward_task_with_type(Task const *task, DT *output_ptr = helperGetTensorPointerWO
( regions[1], task->regions[1], FID_DATA, ctx, runtime); + if (task->index_point.point_data[0] == 0) { + int in_dim0 = input_domain.hi()[0] - input_domain.lo()[0] + 1; + int in_dim1 = input_domain.hi()[1] - input_domain.lo()[1] + 1; + int out_dim0 = output_domain.hi()[0] - output_domain.lo()[0] + 1; + int out_dim1 = output_domain.hi()[1] - output_domain.lo()[1] + 1; + printf("Combine: in=[%i,bz=?/%i] -> out=[%i,bz=?/%i]\n", + in_dim0, in_dim1, + out_dim0, out_dim1); + } + forward_kernel
(input_ptr, output_ptr, output_domain.get_volume()); } diff --git a/src/runtime/ffconst_utils.cc b/src/runtime/ffconst_utils.cc index 9ecd5bd98..c8a76cadc 100644 --- a/src/runtime/ffconst_utils.cc +++ b/src/runtime/ffconst_utils.cc @@ -188,6 +188,8 @@ std::string get_operator_type_name(OperatorType type) { return "Sampling"; case OP_ARGMAX: return "ArgMax"; + case OP_DECODING: + return "Decoding"; // PEFT Ops case OP_LORA: return "Lora Layer"; diff --git a/src/runtime/graph.cc b/src/runtime/graph.cc index fbfbb4d48..079f9e261 100644 --- a/src/runtime/graph.cc +++ b/src/runtime/graph.cc @@ -24,6 +24,7 @@ #include "flexflow/ops/beam_topk.h" #include "flexflow/ops/cast.h" #include "flexflow/ops/concat.h" +#include "flexflow/ops/decoding.h" #include "flexflow/ops/conv_2d.h" #include "flexflow/ops/dropout.h" #include "flexflow/ops/element_binary.h" @@ -3057,6 +3058,10 @@ void FFModel::deserialize_graph_optimal_view( node = ArgMax::deserialize(*this, dez, inputs, num_inputs); break; } + case OP_DECODING: { + node = Decoding::deserialize(*this, dez, inputs, num_inputs); + break; + } case OP_GROUP_BY: { node = Group_by::deserialize(*this, dez, inputs, num_inputs); break; diff --git a/src/runtime/inference_manager.cc b/src/runtime/inference_manager.cc index 14e0d2be8..7f615ca15 100644 --- a/src/runtime/inference_manager.cc +++ b/src/runtime/inference_manager.cc @@ -475,7 +475,7 @@ InferenceResultFuture InferenceManager::inference(FFModel *model, } fm = op->inference(*model, bc, inputs, outputs); } - assert(fm.get_future_map_domain().get_volume() == 1); + assert(fm.get_future_map_domain().get_volume() == model->config.tensor_parallelism_degree); InferenceResultFuture irf = fm.get_future(0); return irf; }; @@ -489,10 +489,13 @@ std::vector InferenceManager::peft_bwd( // Assert that the last operator must be argmax or sampling assert(model->operators[last_op]->op_type == OP_ARGMAX || model->operators[last_op]->op_type == OP_ARG_TOPK || - model->operators[last_op]->op_type == OP_SAMPLING); - last_op -= 1; - while (model->operators[last_op]->op_type == OP_WEIGHT && last_op > 0) { + model->operators[last_op]->op_type == OP_SAMPLING || + model->operators[last_op]->op_type == OP_DECODING); + if (model->operators[last_op]->op_type != OP_DECODING) { last_op -= 1; + while (model->operators[last_op]->op_type == OP_WEIGHT && last_op > 0) { + last_op -= 1; + } } for (int o = last_op; o >= 0; o--) { Op *op = model->operators[o]; diff --git a/src/runtime/model.cc b/src/runtime/model.cc index 959af8779..f7db0c614 100644 --- a/src/runtime/model.cc +++ b/src/runtime/model.cc @@ -27,6 +27,7 @@ #include "flexflow/ops/aggregate_spec.h" #include "flexflow/ops/arg_topk.h" #include "flexflow/ops/argmax.h" +#include "flexflow/ops/decoding.h" #include "flexflow/ops/attention.h" #include "flexflow/ops/batch_matmul.h" #include "flexflow/ops/batch_norm.h" @@ -3364,6 +3365,11 @@ Op *FFModel::create_operator_from_layer( operators.push_back(op); return op; } + case OP_DECODING: { + Op *op = Decoding::create_operator_from_layer(*this, layer, inputs); + operators.push_back(op); + return op; + } case OP_GROUP_BY: { Op *op = Group_by::create_operator_from_layer(*this, layer, inputs); operators.push_back(op); @@ -3432,10 +3438,10 @@ bool FFModel::need_to_add_combine(int layer_idx) const { return false; } } - // argmax/arg_topk not precedent by softmax: add combine before + // argmax/arg_topk not preceded by softmax: add combine before // argmax/arg_topk if (layer_idx == layers.size() - 1 && - (l->op_type == OP_ARG_TOPK || l->op_type == OP_ARGMAX)) { + (l->op_type == OP_ARG_TOPK || l->op_type == OP_ARGMAX)) { // || l->op_type == OP_DECODING auto const &l_prev = layers[layer_idx - 1]; if (l_prev->op_type == OP_SOFTMAX) { return false; @@ -6717,6 +6723,71 @@ void register_flexflow_internal_tasks(Runtime *runtime, registrar); } } + // Decoding task + { + TaskVariantRegistrar registrar(DECODING_INIT_TASK_ID, "Decoding Init"); + registrar.add_constraint(ProcessorConstraint(Processor::TOC_PROC)); + registrar.set_leaf(); + if (pre_register) { + Runtime::preregister_task_variant( + registrar, "Decoding Init Task"); + } else { + if (enable_control_replication) { + registrar.global_registration = false; + } + runtime->register_task_variant(registrar); + } + } + { + TaskVariantRegistrar registrar(DECODING_BEAM_INF_TASK_ID, + "Decoding Beam Inference"); + registrar.add_constraint(ProcessorConstraint(Processor::TOC_PROC)); + registrar.set_leaf(); + if (pre_register) { + Runtime::preregister_task_variant( + registrar, "Decoding Inference Task Beam"); + } else { + if (enable_control_replication) { + registrar.global_registration = false; + } + runtime->register_task_variant(registrar); + } + } + { + TaskVariantRegistrar registrar(DECODING_NORM_INF_TASK_ID, + "Decoding Norm Inference"); + registrar.add_constraint(ProcessorConstraint(Processor::TOC_PROC)); + registrar.set_leaf(); + if (pre_register) { + Runtime::preregister_task_variant( + registrar, "Decoding Inference Task Norm"); + } else { + if (enable_control_replication) { + registrar.global_registration = false; + } + runtime + ->register_task_variant( + registrar); + } + } + { + TaskVariantRegistrar registrar(DECODING_PEFT_BWD_TASK_ID, + "Decoding PEFT Backward"); + registrar.add_constraint(ProcessorConstraint(Processor::TOC_PROC)); + registrar.set_leaf(); + if (pre_register) { + Runtime::preregister_task_variant( + registrar, "Decoding PEFT Backward Task"); + } else { + if (enable_control_replication) { + registrar.global_registration = false; + } + runtime->register_task_variant(registrar); + } + } // Transpose task { TaskVariantRegistrar registrar(TRANSPOSE_INIT_TASK_ID, "Transpose Init"); diff --git a/src/runtime/operator_params.cc b/src/runtime/operator_params.cc index e9feb86eb..00d2d6fe3 100644 --- a/src/runtime/operator_params.cc +++ b/src/runtime/operator_params.cc @@ -5,6 +5,7 @@ #include "flexflow/ops/arg_topk.h" #include "flexflow/ops/argmax.h" #include "flexflow/ops/attention.h" +#include "flexflow/ops/decoding.h" #include "flexflow/ops/batch_matmul.h" #include "flexflow/ops/batch_norm.h" #include "flexflow/ops/beam_topk.h" @@ -150,6 +151,8 @@ tl::optional get_op_parameters(Op const *op) { return ((Sampling *)op)->get_params(); case OP_ARGMAX: return ((ArgMax *)op)->get_params(); + case OP_DECODING: + return ((Decoding *)op)->get_params(); // TODO: implement the get_params() function for the operators below and // uncomment the lines below diff --git a/src/runtime/parallel_tensor.cc b/src/runtime/parallel_tensor.cc index 202983e8f..079083ff3 100644 --- a/src/runtime/parallel_tensor.cc +++ b/src/runtime/parallel_tensor.cc @@ -554,6 +554,10 @@ void ParallelTensorBase::print(std::string const &name) const { for (int i = 0; i < num_dims; i++) { printf("%d ", dims[i].parallel_idx); } + printf("] is_replica_dim["); + for (int i = 0; i < num_dims; i++) { + printf("%d ", dims[i].is_replica_dim ? 1 : 0); + } printf("]\n"); } diff --git a/tests/peft_test.sh b/tests/peft_test.sh index 296097802..6b042ebdd 100755 --- a/tests/peft_test.sh +++ b/tests/peft_test.sh @@ -8,6 +8,9 @@ cleanup() { # Cd into directory holding this script cd "${BASH_SOURCE[0]%/*}/.." +cd build +source ./set_python_envs.sh +cd .. MODEL_NAME=${MODEL_NAME:-"goliaro/llama-3.2-1b-lora"} BASE_MODEL_NAME=${BASE_MODEL_NAME:-"unsloth/Llama-3.2-1B-Instruct"} @@ -43,13 +46,13 @@ mkdir -p ./inference/output export LEGION_BACKTRACE=1 # Download test model -# python ./inference/utils/download_peft_model.py "${MODEL_NAME}" +python ./inference/utils/download_peft_model.py "${MODEL_NAME}" if [ "$FULL_PRECISION" = "true" ]; then full_precision_flag="--use-full-precision"; else full_precision_flag=""; fi if [ "$FUSION" = "true" ]; then fusion_flag="--fusion"; else fusion_flag=""; fi # Run PEFT in Huggingface to get ground truth tensors -# eval python ./tests/peft/hf_finetune.py --peft-model-id "${MODEL_NAME}" --save-peft-tensors "${full_precision_flag}" -lr "${LEARNING_RATE}" +eval python ./tests/peft/hf_finetune.py --peft-model-id "${MODEL_NAME}" --save-peft-tensors "${full_precision_flag}" -lr "${LEARNING_RATE}" # Python test echo "Python test"