Returns the trainer for the job.

Initialize the new model from the pretrained model. This method uses Gopher layer scaling and Phi-style weight tiling as selected by `mode`.

Initialize the new model from the pretrained model. This method uses Gopher layer scaling and Phi-style weight tiling as selected by `mode`.

Main: create C4/pile streaming dataset. Args: args (Namespace): Commandline arguments.

Manage the launching of jobs for each configuration file in the directory.

If shard does not have decompressed data, this function decompresses the shard

Add padding to sequences. Args: inputs: (total_nnz, ...) or (total_nnz,), where total_nnz = number of tokens selected in attention_m

Parse commandline arguments.

Prints fine-tuning results.

Cleans up memory usage left by trainer.

Instantiates the job object and runs it.

Instantiates the job object and runs it.

Functional API to apply a StableAdamW optimization step. See `optimi.StableAdamW` for more details. Args: params: Parameters to upda

Remove padding from input sequences. Args: inputs: (batch, seqlen, ...) or (batch, seqlen) attention_mask: (batch, seqlen),

Takes batches of unpacked, unpadded sequences (seqs) to batches of packed and padded sequences (pseqs). Every input batch must be a

Args: - iterable: an object which generates a fresh iterator on iter() and which implements len()