If distributed is initialized print only on rank 0.

Retrieves a single item from the dataset with the option to only return a portion of the item. get(idx) is the same as [idx] but get(

Search for a given key in a NeoX yaml. normalizes underscores -> hyphens

updates a property value if the key already exists Problem: a previously non-existing property can be added to the class instance wi

Instantiate a PreTrainedBertModel from a pre-trained model file. Download and cache the pre-trained model file if needed.

Return world size for the model parallel group.

Start the timer.

Get the model parallel group the caller rank belongs to.

Initialize affine weight for model parallel. Build the master weight on all processes and scatter the relevant chunk.

Initialize affine weight for model parallel on GPU.

instantiates NeoXArgs while reading values from yml files paths_to_yml_files: list of paths to yml files overwrite_values:

Return my rank for the model parallel group.

Generates a random sample of max_tests length of all possible combinations of values in `params_to_test`. In `params_to_test` you can ei

Stop the timer.

Log a group of timers.

Ensure that numerator is divisible by the denominator and return the division value.

Get the fp32 allreduce flag

True if mp rank == 0

Auxiliary function for checking the values in the checkpoint and setting them.

entry point for deepy.py configuring and consuming command line arguments. We can use `--wandb_group` / `--wandb_team` to overwrite

takes a sequence of parent classes and returns corresponding values (with defaults set)

Initialize tokenizer.

Sets up the model for inference by turning on k/v caching (if specified) and setting `parallel output` of the final layer to false, s

Returns a list containing all weights in a given layer from a model (across MP partitions)

Recursively set attributes on a pytorch module or an iterable of modules. If an assert_type is provided, it will assert that the type of the

Reset timer.

Save a model checkpoint.

Setup model and optimizer.

Updates multiple values in self if the keys already exists

Gather tensors and concatinate along the last dimension.

All-reduce the the input tensor across model parallel group.

Split the tensor along its last dimension and keep the corresponding slice.

Build train/valid/test datasets.

Configures distributed training arguments from local variables set by deepspeed launcher.

Deepspeed launcher needs to pass the arguments for `pretrain_gpt2.py` across to all machines. In order not to have any problems with

convert a NeoX checkpoint to a HF model format. should perform model-parallel merging correctly but only supports features allowed by HF GPT-N

Calculate the elapsed time.

Helper function to evaluate and dump results on screen.

Runs model.forward(model_inputs) We need to create a wrapper for this function because deepspeed pipe parallel modules operate differently t

Generates samples from raw text and returns them in a dictionary. neox_args: NeoXArgs. model: a Megatron model text: either a single

Generate batch from context tokens. Attention mask and position ids are created. Returned tensors will be on CUDA. neox_args: NeoXArgs.

Get the data parallel group the caller rank belongs to.

Set initialize distributed and set autoresume and random seeds. `allow_no_cuda` should not be set unless using megatron for cpu only data proc

Build dataloader given an input dataset.

Transforms the PipelineModule to a plain nn.Sequential module :return:

Single training step.

extracts samples only pertaining to this worker's batch

Support function for get_batch / get_batch pipe (to avoid code repetition)