Reset the memory used for checkpointing.

Save the tokenizer vocabulary and merge files to a directory.

Send tensor to previous rank in pipeline (backward send).

Batched recv from next rank and send to previous rank in pipeline.

Batched send and recv with previous rank in pipeline.

Send tensor to next rank in pipeline (forward send).

Batched send and recv with previous and next ranks in pipeline.

Batched send and recv with next rank in pipeline.

Batched recv from previous rank and send to next rank in pipeline.

See megatron.model.transformer.set_input_tensor()

Set input tensor to be used instead of forward()'s input. When doing pipeline parallelism the input from the previous stage comes fro

r""" Default setting for the pipeline is: `(layer_id + offset) // (layers / pipeline_stage)`. Args: network(Cell) - Repr

r""" Default setting for the pipeline is: `(layer_id + offset) // (layers / pipeline_stage)`. Args: network(Cell) - Repr

Set pipeline model parallel rank.

Set the pipeline model parallel size

Customized load.

Customized load.

Customized load.

Set the virtual pipeline-parallel rank.

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

For easy load.

For easy load.

For easy load.

For easy load.

Use this function to override the state dict for saving checkpoints.

For easy load.

For easy load.

For easy load.

For easy load.

Print loss after each step

step end

task for each process

Get topk

Build train, valid, and test datasets.

Build train, valid, and test datasets.

Truncates a pair of sequences to a maximum sequence length.

Update list of ppl

Forward step.

Dictionary from vocab text token to id token.

Helper function for the cross entropy.

Easy accessory for the DeepSpeed pipeline engine to tie embeddings across stages.

Easy accessory for the DeepSpeed pipeline engine to tie embeddings across stages.

Writes an iterable of dictionaries to jsonl

We only need to zero the model related parameters, i.e., float16_groups & fp32_from_fp32_groups.