A DataProto is a data structure that aims to provide a standard protocol for data exchange between functions. It contains a batch (TensorDict

Sharding manager to support data resharding when using FSDP + Ulysses

This is an in-memory SFTDataset

We assume the dataset contains a column that contains prompts and other information

Note that this trainer runs on the driver process on a single CPU/GPU node.

Define a resource pool specification. Resource pool will be initialized first. Mapping

DataProtoFuture aims to eliminate actual data fetching on driver. By doing so, the driver doesn't have to wait for data so that asynchronous

The implementation of a simple digit completion task. The prompt is a sequence of numbers with fixed difference. The task is to complete the

Used to count mfu during training loop Example: flops_counter = FlopsCounter(config) flops_achieved, flops_promised = flops_

The reward manager.

Adaptive KL controller described in the paper: https://arxiv.org/pdf/1909.08593.pdf

Disable All Megatron Timing with FakeTimers

Fixed KL controller.

LlamaRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla

LlamaRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev

A memory buffer is a contiguous torch tensor that may combine multiple tensors sharing with the underlying memory. It must have a unique type

Note that we do not design MemoryBufferModuleWrapper as an nn.Module due to - It will change the checkpoint name

MegatronWorkerGroup will query each worker of its megatron rank info and store it inside the WorkerGroup so that the dispatcher can use it to

Configuration for optimizer.

Multi-headed attention from 'Attention Is All You Need' paper

Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`] Args: config: LlamaConfig

Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`] Args: config: LlamaConfig

Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`] This model definition supports pip

This worker can be instantiated as a standalone actor or a standalone rollout or a standalone reference policy or a hybrid engine based on th

This worker can be instantiated as a standalone actor or a standalone rollout or a standalone reference policy or a hybrid engine based on th

This class stores a class constructor and the args/kwargs to construct the class. It is used to instantiate the remote class.

the tokenizer property and function name should align with HF's to meet vllm requirement

We assume that - inference engine has tp + dp - actor has tp + pp + dp - the tp between inference engine and actor should be the same

MegatronWorkerGroup will query each worker of its megatron rank info and store it inside the WorkerGroup so that the dispatcher can use it to

Type of precision used. >>> PrecisionType.HALF == 16 True >>> PrecisionType.HALF in (16, "16") True

Note that we only implement the reward model that is subclass of AutoModelForSequenceClassification.

Note that we only implement the reward model that is subclass of AutoModelForTokenClassification.

To create more roles dynamically, you can subclass Role and add new members