Return the virtual pipeline-parallel rank.

Return the virtual pipeline-parallel world size.

Checks whether the generated code is finished.

Check if the current word piece is the starting piece (BERT).

Customized load.

Build the model.

LM logits using word embedding weights.

Main training program. This function will run the followings in the order provided: 1) initialize Megatron. 2) setup model, optim

Set to the initial state (no tracker).

Convert the log_probs to probability

See megatron.model.transformer.set_input_tensor()

Set weight decay coefficient, zero for bias and layernorm, 1e-1 for rest

This function has been mostly taken from huggingface conversational ai code at https://medium.com/huggingface/how-to-build-a-state-of-the-

Get topk

Get topk

Get topk

Memory optimization arguments.

Text generate arguments.

Build an array with length = number-of-epochs * number-of-dcuments. Each index is mapped to a corresponding document.

Build index mappings. We only have to build doc-idx in prompt dataset. Args: name: name of the dataset. data_prefix: prefix o

Build the size on rank 0 and broadcast.

Check that all the keys have the same target data type.

check dynamic format branch

Check the type of inputs.

check parameters including shape_dy, shape_x, shape_var, shape_mean, shape_gamma, dtype and kernel_name Parameters ---------- pa

check parameters including shape_dy, shape_x, shape_var, shape_mean and shape_gamma Parameters ---------- params_map: dict

check if parameter shape_gamma meets the requirements of function Parameters ---------- shape_x: list or tuple shape of x

check if parameter shape_mean meets the requirements of function Parameters ---------- shape_x: list or tuple shape of x

judge case using cube to handle reducesum only supported follow case in Ascend910 and Ascend710: ori_shape: ((batch), m, 1024(768)),

Forward method with activation checkpointing.

get placeholders of data_dy, data_x, data_variance, data_mean and data_gamma Parameters ---------- shapes: dict {"shape_dy":

Calculate the communication group of model parallel dim in one pipeline stage

compute parameters including param_axis, reduce_axis and mean_num Parameters ---------- shape_x: list or tuple shape of x

Divide params into with-weight-decay and without-weight-decay groups. Layernorms and baises will have no weight decay but the rest will.

compute pd_mean according to reduce_axis, pd_xl, pd_var, var_elta_2 and sub_x_mean Parameters ---------- params: dict {"

compute pd_mean according to reduce_axis, pd_xl, pd_var, var_elta_2 and sub_x_mean

compute pd_var according to data_x, data_mean, reduce_axis and pd_xl Parameters ---------- data: dict placeholders after cas

compute front part of pd_var according to data_variance Parameters ---------- data: dict placeholders after cast cast_dt

compute front part of pd_var according to data_variance

compute pd_var according to data_x, data_mean, reduce_axis and pd_xl

compute pd_x, res_for_gamma according to data, params and shape_x `pd_x = pd_xl * np.power((data_variance + EPSLON), (-0.5))`

compute pd_x according to data, params and shape_x

compute pd_xl according to data_dy, data_gamma and shape_x Parameters ---------- data: dict placeholders after cast shap

compute pd_xl according to data_dy, data_gamma and shape_x

get params and data, compute pd_x, pd_gamma, pd_beta. Parameters ---------- data_dy: TVM tensor the placeholder of dy input

get params and data, compute pd_x, pd_gamma, pd_beta.

Calculate the communication group between all pipeline stages

compute pd_x, pd_gamma, pd_beta according to data, params and shape_x Parameters ---------- data: dict placeholders after ca