jinhongyii commented on code in PR #16169: URL: https://github.com/apache/tvm/pull/16169#discussion_r1408499649
########## src/relax/distributed/transform/lower_distir.cc: ########## @@ -0,0 +1,271 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +/*! + * \file tvm/relax/distributed/transform/lower_distir.cc + * \brief Pass for lowering DistIR into Relax + * This pass assumes all the TensorIR functions are in local view, + * so the pass only handles sharding relax tensor shape and + * inserting necessary broadcast and scatter for inputs. + */ + +#include <tvm/relax/attrs/ccl.h> +#include <tvm/relax/distributed/axis_group_graph.h> +#include <tvm/relax/distributed/transform.h> +#include <tvm/relax/expr_functor.h> +#include <tvm/tir/stmt_functor.h> + +#include "../../../tir/schedule/transform.h" +#include "../../op/ccl/ccl.h" +#include "../../op/tensor/manipulate.h" +#include "utils.h" + +namespace tvm { +namespace relax { +namespace distributed { + +class DistIRSharder : public ExprMutator { + public: + static IRModule LowerDistIR(IRModule mod) { return DistIRSharder(mod).Lower(); } + + private: + explicit DistIRSharder(IRModule mod) : ExprMutator(mod) {} + + IRModule Lower() { + auto mod = builder_->GetContextIRModule(); + for (const auto& [gv, base_func] : mod->functions) { + const auto* func_ = base_func.as<FunctionNode>(); + if (func_ == nullptr || !IsDistIRFunc(GetRef<Function>(func_))) { + continue; + } + Function func = RewriteFunction(GetRef<Function>(func_)); + builder_->UpdateFunction(gv, func); + } + return builder_->GetContextIRModule(); + } + + ShapeExpr ShardShape(ShapeExpr orig_shape, DeviceMesh device_mesh, Placement placement) { + ShapeTuple device_mesh_shape = device_mesh->shape; + Array<PrimExpr> new_tensor_shape_value = orig_shape->values; + for (int i = 0; i < device_mesh_shape.size(); i++) { + if (placement->dim_specs[i]->kind == PlacementSpecKind::kSharding) { + int shard_size = device_mesh_shape[i]; + int axis = placement->dim_specs[i]->axis; + new_tensor_shape_value.Set(axis, floordiv(orig_shape->values[axis], shard_size)); + } + } + return ShapeExpr(new_tensor_shape_value); + } + + TensorStructInfo ShardDTensorSinfo(DTensorStructInfo orig_sinfo) { + TensorStructInfo tensor_sinfo = orig_sinfo->tensor_sinfo; + ICHECK(tensor_sinfo->shape); + const auto* orig_shape = tensor_sinfo->shape.as<ShapeExprNode>(); + auto new_tensor_sinfo = make_object<TensorStructInfoNode>(*tensor_sinfo.get()); + new_tensor_sinfo->shape = + ShardShape(GetRef<ShapeExpr>(orig_shape), orig_sinfo->device_mesh, orig_sinfo->placement); + return TensorStructInfo(new_tensor_sinfo); + } + + StructInfo ConvertSinfo(StructInfo orig_sinfo, bool shard_shape) { + if (const auto* dtensor_sinfo = orig_sinfo.as<DTensorStructInfoNode>()) { + if (shard_shape) { + return ShardDTensorSinfo(GetRef<DTensorStructInfo>(dtensor_sinfo)); + } else { + return dtensor_sinfo->tensor_sinfo; + } + } else if (const auto* tuple_sinfo = orig_sinfo.as<TupleStructInfoNode>()) { + Array<StructInfo> new_fields; + for (const auto& field_sinfo : tuple_sinfo->fields) { + if (const auto* dtensor_sinfo = field_sinfo.as<DTensorStructInfoNode>()) { + if (shard_shape) { + new_fields.push_back(ShardDTensorSinfo(GetRef<DTensorStructInfo>(dtensor_sinfo))); + } else { + new_fields.push_back(dtensor_sinfo->tensor_sinfo); + } + } else { + new_fields.push_back(field_sinfo); + } + } + return TupleStructInfo(new_fields); + } else { + return orig_sinfo; + } + } + + Expr ShardInputParamTensorAndConstant(Expr input) { + ICHECK(input->struct_info_); + StructInfo old_sinfo = GetStructInfo(input); + StructInfo new_sinfo = ConvertSinfo(old_sinfo, false); + if (const auto* var = input.as<VarNode>()) { + Var new_param(var->name_hint(), new_sinfo); + return new_param; + } else if (const auto* constant = input.as<ConstantNode>()) { + for (const auto& spec : Downcast<DTensorStructInfo>(old_sinfo)->placement->dim_specs) { + ICHECK(spec->kind == PlacementSpecKind::kReplica); + } + Constant new_constant(constant->data, new_sinfo); + return new_constant; + } else { + LOG(FATAL) << "Cannot shard tensor which is not Var or Constant: " << input; + throw; + } + } + + void EmitBroadcastOrScatter(Expr old_expr, Expr new_expr, DTensorStructInfo dtensor_sinfo) { + // FIXME: this is a hack that only works for 1d device mesh + ICHECK(dtensor_sinfo->device_mesh->shape.size() == 1); + PlacementSpec sharding_spec = dtensor_sinfo->placement->dim_specs[0]; + if (sharding_spec->kind == PlacementSpecKind::kReplica) { + Var new_var = builder_->Emit(broadcast_from_worker0(new_expr)); + if (const auto* var = old_expr.as<VarNode>()) { + var_remap_[var->vid] = new_var; + } else { + tuple_getitem_remap_[Downcast<TupleGetItem>(old_expr)] = new_var; + } + } else if (sharding_spec->kind == PlacementSpecKind::kSharding) { + Var scatter_var = builder_->Emit(scatter_from_worker0( + new_expr, dtensor_sinfo->device_mesh->shape[0], sharding_spec->axis)); + if (const auto* var = old_expr.as<VarNode>()) { + var_remap_[var->vid] = scatter_var; + } else { + tuple_getitem_remap_[Downcast<TupleGetItem>(old_expr)] = scatter_var; + } + } else { + LOG(FATAL) << "Unsupported placement spec"; + } + } + + void InputPreprocessing() { + for (int i = 0; i < static_cast<int>(func_->params.size()); i++) { + Var param = func_->params[i]; + if (const auto* dtensor_sinfo = GetStructInfoAs<DTensorStructInfoNode>(param)) { + EmitBroadcastOrScatter(param, new_params_[i], GetRef<DTensorStructInfo>(dtensor_sinfo)); + } else if (const auto* tuple_sinfo = GetStructInfoAs<TupleStructInfoNode>(param)) { + for (int j = 0; j < static_cast<int>(tuple_sinfo->fields.size()); j++) { + if (const auto* dtensor_sinfo = tuple_sinfo->fields[j].as<DTensorStructInfoNode>()) { Review Comment: It's possible that the input contains a non-DTensorStructInfoNode element. For example, kv cache is object. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
