Initialize Shampoo preconditioners and inverse preconditioners.

Compute matrix square root using coupled inverse Newton iteration. alpha <- -1 / p X <- 1/c * I M <- 1/c^p * A repeat

Compute matrix (inverse) root using eigendecomposition of symmetric positive (semi-)definite matrix. A = Q L Q^T => A^{1/r} = Q L^{1/r} Q

Save streaming dataset to a fixed mapping-style database.

Updates preconditioners. Note: If using L2-regularization/weight decay, it is computed within this function and therefore should not be

tuformer attention in batch first implementation (hopefully)

tuformer attention in batch first implementation (hopefully)

These measurements are deliberately wall-clock based.

Early termination based on terrible loss.

unbiased cov test.

See the config file for details on what is possible.

construct model from given configuration. Only works if this arch exists on the hub.

See the config file for details on what is possible.

See the config file for details on what is possible.

See the config file for details on what is possible.

Experimental.

Create an LMDB database from the given pytorch dataset. https://github.com/Lyken17/Efficient-PyTorch/blob/master/tools/folder2lmdb.py Remove

"Seamlessly" run exact deduplication as in Lee et al.

A quick and dirty communication through NCCL. Should not be a major burden.

Straight-up from https://stackoverflow.com/a/6027615/3775820.

Inputs of shape [B, S, H]. Implementation directly based on FLASH pseudocode (see paper appendix)

compute the Hessian vector product with a random vector v, at the current gradient point, i.e., compute the gradient of <gradsH,v>.

Run data-based initialization search as described in Zhu et al., "GradInit: Learning to Initialize Neural Networks for Stable and Efficient Tr

Initialize deepspeed. Module is imported lazily here.

initialize a torch engine.

Inner product similar to description in Bollapragada,Byrd,Nocedal, "Adaptive Sampling Strategies for Stochastic Optimization". This is only a

Computes matrix root inverse. Args: A (Tensor): Square matrix of interest. root (int): Root of interest. Any natural number.

Reshapes tensor by merging small dimensions. Args: tensor_size (Tensor or List[int]): The shape of the tensor. threshold (int): T

Concatenates multiple tensors to form single tensor across multiple dimensions. Args: split_grad (List[Tensor]): List of gradient chunks.

Sohams version.

A lot of loading and preprocessing.

Analogous to save_final_model, but save model to hugginface hub.

Some dataset "improvements". These are optional filtering or normalization rules that are only applied to the pretraining corpus. This separates t

Tokens in each microbatch step.

This checkpoint can be used for downstream tasks. The default behavior is to save this checkpoint to a checkpoints folder under base_directory

Path, identifier and additional client state. This checkpoint can be used to resume training. The default behavior is to save this checkpoint

Optionally implement linear batch size ramp-ups.

Optionally implement linear sequence lengths curriculum.

Switch pytorch into a deterministic computation mode.

.

Assume huggingface data is B S

Allow dynamic modifications of batch size.

Decide and print GPU / CPU / hostname info. Generate local distributed setting if running in distr. mode. Set all required and interesting enviro

Move batch of data into device memory.

Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original.

Update exponential moving average in parameters and buffers.

Update sample statistics based on welford accumulation. At any step variance can be finalized via running_variance / count

Call this only if all other attributes are exhausted.

Call this only if all attributes of Subset are exhausted.

This is only called if all attributes of Subset are exhausted.

Get from database. This is either unordered or cursor access for now. Future: Write this class as a proper https://pytorch.org/docs/stable/da

Get sample, target from cache.