Delete folder olmo_bitnet_1b with huggingface_hub

olmo_bitnet_1b/.gitattributes

@@ -1,35 +0,0 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

olmo_bitnet_1b/README.md

@@ -1,38 +0,0 @@
----
-license: apache-2.0
-datasets:
-- allenai/dolma
----
-# OLMo-Bitnet-1B
-
-OLMo-Bitnet-1B is a 1B parameter model trained using the method described in [The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits](https://arxiv.org/abs/2402.17764).
-
-It was trained on the first 60B tokens of the [Dolma](https://huggingface.co/datasets/allenai/dolma) dataset, so it is merely a research proof-of-concept to test out the methodolgy.
-
-A separate training run was run with the exact same hyperparameters, but using standard fp16 weights.
-The comparison can be found in [this wandb report](https://api.wandb.ai/links/emozilla/evltqiv7).
-
-
-![image/png](https://cdn-uploads.huggingface.co/production/uploads/6317aade83d8d2fd903192d9/NAw-hyWJl5ihVsAPqz3Xe.png)
-
-Sample inference code
-
-```sh
-pip install ai2-olmo
-```
-
-```python
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, TextStreamer
-
-tokenizer = AutoTokenizer.from_pretrained("NousResearch/OLMo-Bitnet-1B")
-model = AutoModelForCausalLM.from_pretrained("NousResearch/OLMo-Bitnet-1B",
-    torch_dtype=torch.bfloat16, trust_remote_code=True, device_map="auto")
-
-streamer = TextStreamer(tokenizer)
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, pad_token_id=tokenizer.eos_token_id,
-    temperature=0.8, repetition_penalty=1.1, do_sample=True,streamer=streamer)
-pipe("The capitol of Paris is",  max_new_tokens=256)
-```
-
-Training was performed using [OLMo](https://github.com/allenai/OLMo).

olmo_bitnet_1b/aliases.py

@@ -1,7 +0,0 @@
-from os import PathLike
-from typing import Union
-
-__all__ = ["PathOrStr"]
-
-
-PathOrStr = Union[str, PathLike]

olmo_bitnet_1b/beam_search.py

@@ -1,1078 +0,0 @@
-"""
-This is a self-contained and flexible beam search implementation adapted from
-AllenNLP's beam search: https://github.com/allenai/allennlp/blob/main/allennlp/nn/beam_search.py
-"""
-
-import copy
-import warnings
-from abc import abstractmethod
-from inspect import signature
-from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, cast
-
-import torch
-
-__all__ = [
-    "Sampler",
-    "DeterministicSampler",
-    "MultinomialSampler",
-    "TopKSampler",
-    "TopPSampler",
-    "GumbelSampler",
-    "FinalSequenceScorer",
-    "SequenceLogProbabilityScorer",
-    "LengthNormalizedSequenceLogProbabilityScorer",
-    "Constraint",
-    "RepeatedNGramBlockingConstraint",
-    "BeamSearch",
-]
-
-StateType = Dict[str, torch.Tensor]
-StepFunctionTypeWithTimestep = Callable[[torch.Tensor, StateType, int], Tuple[torch.Tensor, StateType]]
-StepFunctionTypeNoTimestep = Callable[[torch.Tensor, StateType], Tuple[torch.Tensor, StateType]]
-
-StepFunctionType = TypeVar("StepFunctionType", StepFunctionTypeWithTimestep, StepFunctionTypeNoTimestep)
-"""
-The type of step function that can be passed to [`BeamSearch.search`](#search).
-
-This can either be [`StepFunctionTypeWithTimestep`](#stepfunctiontypewithtimestep)
-or [`StepFunctionTypeNoTimestep`](#stepfunctiontypenotimestep).
-"""
-
-ConstraintStateType = List[List[Dict[str, Any]]]
-
-
-class Sampler:
-    """
-    An abstract class that can be used to sample candidates (either nodes or beams)
-    within `BeamSearch`.
-
-    A `Sampler` just has three methods, `init_state()`, `sample_nodes()` and `sample_beams()`.
-
-    `init_state()` takes three arguments:
-
-    - a tensor of starting log probs with shape `(batch_size,, num_classes)`,
-    - the batch size, an int,
-    - and the number of classes, also an int.
-
-    It returns a state dictionary with any state tensors needed for subsequent
-    calls to `sample_nodes()` and `sample_beams()`.
-
-    By default this method just returns an empty dictionary.
-
-    Both `sample_nodes()` and `sample_beams()` should take three arguments:
-
-    - tensor of normalized log probabilities with shape `(batch_size, num_examples)`,
-    - an integer representing the number of samples to take for each example in the batch,
-    - and a state dictionary which could contain any tensors needed for the `Sampler` to keep
-      track of state.
-
-    For `sample_nodes()`, `num_examples = num_classes`, but for `sample_beams`,
-    `num_examples = beam_size * per_node_beam_size`.
-
-    The return value should be a tuple containing:
-
-    - a tensor of log probabilities of the sampled examples with shape `(batch_size, num_samples)`,
-    - a tensor of indices of the sampled examples with shape `(batch_size, num_samples)`,
-    - and the updated state dictionary.
-
-    A default implementation of `sample_beams` is provided, which just deterministically
-    picks the `k` examples with highest log probability.
-    """
-
-    def init_state(
-        self, start_class_log_probabilities: torch.Tensor, batch_size: int, num_classes: int
-    ) -> StateType:
-        del start_class_log_probabilities, batch_size, num_classes
-        return {}
-
-    @abstractmethod
-    def sample_nodes(
-        self, log_probs: torch.Tensor, per_node_beam_size: int, state: StateType
-    ) -> Tuple[torch.Tensor, torch.Tensor, StateType]:
-        raise NotImplementedError
-
-    def sample_beams(
-        self, log_probs: torch.Tensor, beam_size: int, state: StateType
-    ) -> Tuple[torch.Tensor, torch.Tensor, StateType]:
-        del state
-        selected_log_probs, selected_indices = torch.topk(log_probs, beam_size, dim=-1)
-        return selected_log_probs, selected_indices, {}
-
-
-class DeterministicSampler(Sampler):
-    """
-    A `Sampler` that just deterministically returns the `k` nodes or beams with highest
-    log probability.
-    """
-
-    def sample_nodes(
-        self, log_probs: torch.Tensor, per_node_beam_size: int, state: StateType
-    ) -> Tuple[torch.Tensor, torch.Tensor, StateType]:
-        del state
-        selected_log_probs, selected_indices = torch.topk(log_probs, per_node_beam_size, dim=-1)
-        return selected_log_probs, selected_indices, {}
-
-
-class MultinomialSampler(Sampler):
-    """
-    A `Sampler` which samples nodes from the given multinomial distribution. Beams are sampled
-    in the default, non-deterministic way.
-
-    :param temperature: A `temperature` below 1.0 produces a sharper probability distribution and a `temperature`
-        above 1.0 produces a flatter probability distribution.
-    :param with_replacement: Whether to sample with replacement.
-
-    """
-
-    def __init__(
-        self,
-        temperature: float = 1.0,
-        with_replacement: bool = False,
-    ) -> None:
-        self.temperature = temperature
-        self.with_replacement = with_replacement
-
-    def sample_nodes(
-        self, log_probs: torch.Tensor, per_node_beam_size: int, state: StateType
-    ) -> Tuple[torch.Tensor, torch.Tensor, StateType]:
-        if self.temperature != 1.0:
-            _probabilities = torch.nn.functional.softmax(log_probs / self.temperature, dim=-1)
-        else:
-            _probabilities = log_probs.exp()
-
-        selected_indices = torch.multinomial(_probabilities, per_node_beam_size, replacement=self.with_replacement)
-
-        return torch.gather(log_probs, 1, selected_indices), selected_indices, state
-
-
-class TopKSampler(Sampler):
-    """
-    A `Sampler` which redistributes the probability mass function for nodes among the
-    top `k` choices, then samples from that subset after re-normalizing the probabilities.
-
-    Beams are sampled in the default, deterministic way.
-
-    :param k: The number of top choices to be selected from.
-    :param temperature: A `temperature` below 1.0 produces a sharper probability distribution and a `temperature`
-        above 1.0 produces a flatter probability distribution.
-    :param with_replacement: If set to `True`, samples will be selected with replacement from the top k choices.
-    """
-
-    def __init__(
-        self,
-        k: int = 1,
-        temperature: float = 1.0,
-        with_replacement: bool = False,
-    ):
-        self.k = k
-        self.temperature = temperature or 1.0
-        self.with_replacement = with_replacement
-
-    def sample_nodes(
-        self, log_probs: torch.Tensor, per_node_beam_size: int, state: StateType
-    ) -> Tuple[torch.Tensor, torch.Tensor, StateType]:
-        if not per_node_beam_size <= self.k <= log_probs.size()[1]:
-            raise ValueError(
-                "k must be a postive integer no less than per_node_beam_size and no greater than vocabulary size"
-            )
-
-        # shape (both): (batch_size, k)
-        top_k_log_probs, top_k_indices = log_probs.topk(self.k, dim=-1)
-
-        # Apply temperature if necessary.
-        # shape: (batch_size, k)
-        if self.temperature != 1.0:
-            top_k_log_probs = top_k_log_probs / self.temperature
-
-        # Re-normalize the subset.
-        # shape: (batch_size, k)
-        normalized_top_k_probs = torch.nn.functional.softmax(top_k_log_probs, dim=-1)
-
-        # Sample from the re-normalized subset.
-        # NOTE: These indices are not indices into `log_probs`, they are indices into `top_k_log_probs`.
-        # shape: (batch_size, per_node_beam_size)
-        sampled_indices = torch.multinomial(
-            normalized_top_k_probs, per_node_beam_size, replacement=self.with_replacement
-        )
-
-        # Convert `sampled_indices` back to indices in the original `log_probs` tensor.
-        # shape: (batch_size, per_node_beam_size)
-        indices = top_k_indices.gather(-1, sampled_indices)
-
-        return log_probs.gather(1, indices), indices, state
-
-
-class TopPSampler(Sampler):
-    """
-    A `Sampler` which redistributes the probability mass function for nodes among
-    the top choices with a cumulative probability of at least `p`, then samples from that subset
-    after re-normalizing the probabilities.
-
-    Beams are sampled in the default, deterministic way.
-
-    :param p:
-        The cumulative probability cutoff threshold. A higher value of `p` will result in more possible
-        examples to sample from. If `with_replacement` is `False` and the number of possible samples is
-        insufficient to sample without replacement from when calling `sample_nodes`, then the top
-        `per_node_beam_size` examples will be chosen.
-    :param temperature:
-        A `temperature` below 1.0 produces a sharper probability distribution and a `temperature`
-        above 1.0 produces a flatter probability distribution.
-    :param with_replacement:
-        If set to `True`, samples will be selected with replacement from the top choices.
-
-    """
-
-    def __init__(
-        self,
-        p: float = 0.9,
-        temperature: float = 1.0,
-        with_replacement: bool = False,
-    ):
-        if p < 0.0 or p > 1.0:
-            raise ValueError("p must be a positive float no greater than 1.0")
-        self.p = p
-        self.temperature = temperature or 1.0
-        self.with_replacement = with_replacement
-
-    def sample_nodes(
-        self, log_probs: torch.Tensor, per_node_beam_size: int, state: StateType
-    ) -> Tuple[torch.Tensor, torch.Tensor, StateType]:
-        if not per_node_beam_size <= log_probs.size()[1]:
-            raise ValueError("per_node_beam_size cannot be greater than vocabulary size")
-
-        # First apply temperature coefficient:
-        if self.temperature != 1.0:
-            _log_probs = torch.nn.functional.log_softmax(log_probs / self.temperature, dim=-1)
-        else:
-            _log_probs = log_probs
-
-        # Sort the probabilities in descending order to then find cumulative sum
-        log_probs_descending, sorting_indices = torch.sort(_log_probs, descending=True)
-
-        # shape: (batch_size, num_classes)
-        probabilities_descending = log_probs_descending.exp()
-        probabilities_summed = torch.cumsum(probabilities_descending, dim=-1)
-
-        # Create a mask for filtering out probabilities that don't make the top `p`.
-        # shape: (batch_size, num_classes)
-        exclusion_mask = probabilities_summed >= self.p
-
-        # We want to include the first index where probabilities_summed >= p, so we shift over one.
-        exclusion_mask[..., 1:] = exclusion_mask[..., :-1].clone()
-        exclusion_mask[..., 0] = False
-
-        # Make sure there's at least `per_node_beam_size` options to be selected.
-        if not self.with_replacement:
-            exclusion_mask[..., :per_node_beam_size] = False
-
-        log_probs_descending[exclusion_mask] = torch.finfo(log_probs.dtype).min
-
-        # Now re-normalized the included log probs.
-        # shape: (batch_size, num_classes)
-        filtered_probabilities = torch.nn.functional.softmax(log_probs_descending, dim=-1)
-
-        # Sample from the re-normalized subset.
-        # NOTE: These indices are not indices into `log_probs`, they are indices into `log_probs_descending`.
-        # shape: (batch_size, per_node_beam_size)
-        sampled_indices = torch.multinomial(
-            filtered_probabilities, per_node_beam_size, replacement=self.with_replacement
-        )
-
-        # Convert `sampled_indices` back to indices in the original `log_probs` tensor.
-        # shape: (batch_size, per_node_beam_size)
-        selected_indices = sorting_indices.gather(-1, sampled_indices)
-
-        # Return (selected log probabilities, selected classes)
-        # shape: (len(log_probs),1) , (len(log_probs), 1)
-        return torch.gather(log_probs, 1, selected_indices), selected_indices, state
-
-
-class GumbelSampler(Sampler):
-    """
-    A `Sampler` which uses the Gumbel-Top-K trick to sample without replacement. See
-    [*Stochastic Beams and Where to Find Them: The Gumbel-Top-k Trick for Sampling
-    Sequences Without Replacement*, W Kool, H Van Hoof and M Welling, 2010]
-    (https://api.semanticscholar.org/CorpusID:76662039).
-
-    :param temperature: A `temperature` below 1.0 produces a sharper probability distribution and a `temperature`
-        above 1.0 produces a flatter probability distribution.
-    """
-
-    def __init__(self, temperature: float = 1.0):
-        self.temperature = temperature
-
-    def init_state(
-        self, start_class_log_probabilities: torch.Tensor, batch_size: int, num_classes: int
-    ) -> StateType:
-        # shape: (batch_size, num_classes)
-        zeros = start_class_log_probabilities.new_zeros((batch_size, num_classes))
-
-        # shape: (batch_size, num_classes)
-        G_phi_S = self.gumbel_with_max(start_class_log_probabilities, zeros)
-
-        return {"G_phi_S": G_phi_S}
-
-    def sample_nodes(
-        self,
-        log_probs: torch.Tensor,
-        per_node_beam_size: int,
-        state: StateType,
-    ) -> Tuple[torch.Tensor, torch.Tensor, StateType]:
-        # First apply temperature coefficient:
-        # shape: (batch_size * beam_size, num_classes)
-        if self.temperature != 1.0:
-            _log_probs = torch.nn.functional.log_softmax(log_probs / self.temperature, dim=-1)
-        else:
-            _log_probs = log_probs
-
-        # shape: (group_size,)
-        phi_S = state["phi_S"]
-
-        # shape: (group_size, num_classes)
-        phi_S = phi_S.unsqueeze(-1).expand_as(_log_probs)
-
-        # shape: (group_size, num_classes)
-        phi_S_new = phi_S + _log_probs
-
-        # shape: (group_size, 1)
-        G_phi_S = state["G_phi_S"].unsqueeze(-1)
-
-        # shape: (group_size, num_classes)
-        G_phi_S_new = self.gumbel_with_max(phi_S_new, G_phi_S)
-
-        # Replace NaNs with very negative number.
-        # shape: (group_size, num_classes)
-        #  G_phi_S_new[G_phi_S_new.isnan()] = torch.finfo(G_phi_S_new.dtype).min
-
-        # shape (both): (group_size, per_node_beam_size)
-        top_G_phi_S_new, top_indices = torch.topk(G_phi_S_new, per_node_beam_size, dim=-1)
-
-        # shape: (group_size, per_node_beam_size)
-        top_log_probs = log_probs.gather(1, top_indices)
-
-        return top_log_probs, top_indices, {"G_phi_S": top_G_phi_S_new}
-
-    def sample_beams(
-        self,
-        log_probs: torch.Tensor,
-        beam_size: int,
-        state: StateType,
-    ) -> Tuple[torch.Tensor, torch.Tensor, StateType]:
-        """
-        Returns the beams with the highest perturbed log probabilities.
-        """
-        # shape (log_probs): (batch_size, beam_size * per_node_beam_size)
-
-        batch_size = log_probs.size()[0]
-
-        # shape: (batch_size * beam_size, per_node_beam_size)
-        G_phi_S = state["G_phi_S"]
-
-        # shape: (batch_size, beam_size * per_node_beam_size)
-        G_phi_S = G_phi_S.reshape_as(log_probs)
-
-        # shape (both): (batch_size, beam_size)
-        G_phi_S_new, selected_indices = torch.topk(G_phi_S, beam_size, dim=-1)
-
-        # shape: (batch_size, beam_size)
-        selected_log_probs = log_probs.gather(1, selected_indices)
-
-        # Now sort the selected beams by their true log prob.
-        # shape (all): (batch_size, beam_size)
-        selected_log_probs, sort_indices = selected_log_probs.sort(dim=-1, descending=True)
-        selected_indices = selected_indices.gather(1, sort_indices)
-        G_phi_S_new = G_phi_S_new.gather(1, sort_indices)
-
-        # shape: (batch_size * beam_size,)
-        G_phi_S_new = G_phi_S_new.reshape(batch_size * beam_size)
-
-        # shape: (batch_size * beam_size,)
-        phi_S = selected_log_probs.reshape(batch_size * beam_size)
-
-        return selected_log_probs, selected_indices, {"G_phi_S": G_phi_S_new, "phi_S": phi_S}
-
-    def gumbel(self, phi) -> torch.Tensor:
-        """
-        Sample `Gumbel(phi)`.
-
-        `phi` should have shape `(batch_size, num_classes)`.
-        """
-        return -torch.log(-torch.log(torch.rand_like(phi))) + phi
-
-    def gumbel_with_max(self, phi, T) -> torch.Tensor:
-        """
-        Sample `Gumbel(phi)` conditioned on the maximum value being equal to `T`.
-
-        `phi` should have shape `(batch_size, num_classes)` and `T` should have
-        shape `(batch_size, 1)`.
-        """
-        # Shape: (batch_size, num_classes)
-        G_phi = self.gumbel(phi)
-
-        # Now we find the maximum from these samples.
-        # Shape: (batch_size, )
-        Z, _ = G_phi.max(dim=-1)
-
-        # Shape: (batch_size, num_classes)
-        v = T - G_phi + torch.log1p(-torch.exp(G_phi - Z.unsqueeze(-1)))
-
-        # Shape: (batch_size, num_classes)
-        return T - torch.nn.functional.relu(v) - torch.log1p(torch.exp(-v.abs()))
-
-
-class FinalSequenceScorer:
-    """
-    An abstract class that can be used to score the final generated sequences found
-    by beam search. Given the predicted sequences and the corresponding log probabilities of
-    those sequences, the class calculates and returns the final score of the sequences.
-
-    The default implementation scores the sequences using the sum of the log probabilities of
-    the sequence, which is passed as input.
-    """
-
-    @abstractmethod
-    def score(self, predictions: torch.Tensor, log_probabilities: torch.Tensor, end_index: int) -> torch.Tensor:
-        """
-        Score the final predictions found by beam search.
-        Returns a tensor of the final sequence scores of shape `(batch_size, beam_size)`.
-
-        :param predictions: A tensor containing the initial predictions with shape `(batch_size, beam_size, max_steps)`.
-        :param log_probabilities: A tensor containing the log probabilities of the sequence, defined as the sum
-            of the log probabilities per token, with shape `(batch_size, beam_size)`.
-        :param end_index: The index of the end symbol.
-
-        """
-        raise NotImplementedError
-
-
-class SequenceLogProbabilityScorer(FinalSequenceScorer):
-    """
-    A :class:`FinalSequenceScorer` which scores the sequences by the sum of the log probabilities
-    across the sequence's tokens.
-    """
-
-    def score(self, predictions: torch.Tensor, log_probabilities: torch.Tensor, end_index: int) -> torch.Tensor:
-        del predictions, end_index
-        # The sum of the sequence log probabilities is the input parameter, so just
-        # return it.
-        return log_probabilities
-
-
-class LengthNormalizedSequenceLogProbabilityScorer(FinalSequenceScorer):
-    """
-    A :class:`FinalSequenceScorer` which scores the sequences by the average log probability of the
-    tokens in the sequence. It optionally includes a length penalty which promotes
-    or demotes sequences based on their lengths. The final score for a sequence will
-    be `(sequence_log_probability) / (sequence_length ** length_penalty)`. The sequence length
-    here includes the end token.
-
-    :param length_penalty: The length penalty to use. A value of 1.0 means no length penalty is used.
-        A value > 1.0 favors longer sequences, and < 1.0 favors shorter sequences.
-    """
-
-    def __init__(self, length_penalty: float = 1.0):
-        super().__init__()
-        self.length_penalty = length_penalty
-
-    def score(self, predictions: torch.Tensor, log_probabilities: torch.Tensor, end_index: int) -> torch.Tensor:
-        # shape: (batch_size, beam_size)
-        lengths = (predictions != end_index).long().sum(dim=2)
-
-        # If the sequence ended during beam search, the `log_probabilities` will include
-        # the transition to the end token. Therefore, in such situations, `lengths` is
-        # actually off by 1. This corrects for that.
-        # shape: (batch_size, beam_size)
-        is_end_token = predictions[:, :, -1] == end_index
-        lengths += is_end_token.long()
-
-        # shape: (batch_size, beam_size)
-        average_log_probs = log_probabilities / (lengths**self.length_penalty)
-        return average_log_probs
-
-
-class Constraint:
-    """
-    An abstract class that can be used to enforce constraints on the output predictions
-    by manipulating the class log probabilities during beam search.
-
-    A `Constraint` just has three methods that need to be implemented by subclasses:
-    `init_state()`, `apply()` and `_update_state()`.
-
-    `init_state()` takes one argument:
-
-    - the batch size, an int
-
-    It returns a constraint state, which is a nested list of dictionaries, with any state needed for subsequent
-    calls to `apply()` and `update_state()`. The length of the outer list should be equal to `batch_size`.
-    Each inner list should be of length 1.
-
-    `apply()` takes two arguments:
-
-    - the constraint state, which is a nested list of dictionaries. The length of the outer list is `batch_size`
-    and the length of each inner list is `beam_size` except on the first time `apply()` is called when it is 1.
-    - `class_log_probabilities`, a tensor of shape `(batch_size, beam_size, num_classes)` that contains the
-    log probabilities for the classes during search. The first time `apply()` is called, `beam_size = 1`.
-
-    The `apply()` method should return new `class_log_probabilities` that enforce the constraint
-    for this step of beam search. For instance, it may prevent a specific class from being selected by setting
-    the corresponding log probability to a negligible value such as `float("-inf")` or
-    `torch.finfo(class_log_probabilities.dtype).min`.
-
-    `_update_state()` takes two arguments:
-
-    - the copied parent constraint state, which is a nested list of dictionaries. `state[i][j]` contains the
-    copied state for the parent of `last_prediction[i, j]`. It is unique to that batch and beam, so it can be
-    directly edited in-place without affecting the others.
-    - last_prediction, a tensor of shape `(batch_size, beam_size)` containing the predictions from the last
-    step of beam search.
-
-    The `_update_state()` function should return a new constraint state, a nested list of dictionaries of
-    length `batch_size` and inner list of length `beam_size`, one for each of the predictions in `last_prediction`.
-
-    """
-
-    @abstractmethod
-    def init_state(
-        self,
-        batch_size: int,
-    ) -> ConstraintStateType:
-        raise NotImplementedError
-
-    @abstractmethod
-    def apply(
-        self,
-        state: ConstraintStateType,
-        class_log_probabilities: torch.Tensor,
-    ) -> torch.Tensor:
-        raise NotImplementedError
-
-    @staticmethod
-    def _copy_state(
-        state: ConstraintStateType,
-        batch_size: int,
-        beam_size: int,
-        last_backpointer: Optional[torch.Tensor] = None,
-    ) -> ConstraintStateType:
-        """
-        Copies the `state` . This method copies the data in `state` using `copy.deepcopy()`. If this
-        is not appropriate for your constraint, you will need to implement the copying yourself.
-        """
-        new_state = []
-        for i in range(batch_size):
-            batch_state = []
-            for j in range(beam_size):
-                if last_backpointer is None:
-                    # This is the first prediction, so the backpointer is 0
-                    backpointer = 0
-                else:
-                    backpointer = last_backpointer[i, j].item()
-                batch_state.append(copy.deepcopy(state[i][backpointer]))  # type: ignore
-            new_state.append(batch_state)
-        return new_state
-
-    def update_state(
-        self,
-        state: ConstraintStateType,
-        last_prediction: torch.Tensor,
-        last_backpointer: Optional[torch.Tensor] = None,
-    ) -> ConstraintStateType:
-        batch_size, beam_size = last_prediction.size()
-        new_state = self._copy_state(state, batch_size, beam_size, last_backpointer)
-        return self._update_state(new_state, last_prediction)
-
-    @abstractmethod
-    def _update_state(
-        self,
-        state: ConstraintStateType,
-        last_prediction: torch.Tensor,
-    ) -> ConstraintStateType:
-        raise NotImplementedError
-
-
-class RepeatedNGramBlockingConstraint(Constraint):
-    def __init__(self, ngram_size: int, **kwargs) -> None:
-        super().__init__(**kwargs)
-        self.ngram_size = ngram_size
-
-    def init_state(
-        self,
-        batch_size: int,
-    ) -> ConstraintStateType:
-        return [[{"seen_ngrams": {}, "current_prefix": []}] for _ in range(batch_size)]
-
-    def apply(
-        self,
-        state: ConstraintStateType,
-        class_log_probabilities: torch.Tensor,
-    ) -> torch.Tensor:
-        for i, batch in enumerate(state):
-            for j, beam in enumerate(batch):
-                current_prefix = tuple(beam["current_prefix"])
-                seen_ngrams = beam["seen_ngrams"]
-                try:
-                    disallowed_indices = seen_ngrams[current_prefix]
-                    class_log_probabilities[i, j, disallowed_indices] = torch.finfo(
-                        class_log_probabilities.dtype
-                    ).min
-                except KeyError:
-                    # We have not seen this prefix before, so there is no index
-                    # that needs to be blocked
-                    pass
-        return class_log_probabilities
-
-    def _update_state(
-        self,
-        state: ConstraintStateType,
-        last_prediction: torch.Tensor,
-    ) -> ConstraintStateType:
-        for i, batch in enumerate(state):
-            for j, beam in enumerate(batch):
-                prediction = last_prediction[i, j].item()
-                prefix = beam["current_prefix"]
-                seen_ngrams = beam["seen_ngrams"]
-
-                if len(prefix) == self.ngram_size - 1:
-                    # This is a new ngram that we have to remember
-                    if tuple(prefix) not in seen_ngrams:
-                        seen_ngrams[tuple(prefix)] = []
-                    seen_ngrams[tuple(prefix)].append(prediction)
-
-                # Create the new prefix, removing the oldest index if the prefix
-                # is too long
-                prefix.append(prediction)
-                if len(prefix) == self.ngram_size:
-                    prefix.pop(0)
-        return state
-
-
-class BeamSearch:
-    """
-    Implements the beam search algorithm for decoding the most likely sequences.
-
-    :param end_index: The index of the "stop" or "end" token in the vocabulary. Usually the EOS token ID.
-
-    :param max_steps: The maximum number of decoding steps to take, i.e. the maximum length
-        of the predicted sequences.
-
-    :param beam_size: The width of the beam used.
-
-    :param per_node_beam_size: The maximum number of candidates to consider per node, at each step in the search.
-        If not given, this just defaults to `beam_size`. Setting this parameter
-        to a number smaller than `beam_size` may give better results, as it can introduce
-        more diversity into the search. See
-        [*Beam Search Strategies for Neural Machine Translation*, Freitag and Al-Onaizan, 2017]
-        (https://api.semanticscholar.org/CorpusID:2229477).
-
-    :param sampler: An optional `Sampler` which is used to pick next candidate nodes and beams.
-        If not specified, `DeterministicSampler` will be used, which just takes the
-        `per_node_beam_size` most likely nodes and the `beam_size` most likely beams.
-
-        Using the [`GumbelSampler`](#gumbelsampler), on the other hand, will give you
-        [Stochastic Beam Search](https://api.semanticscholar.org/CorpusID:76662039).
-
-    :param min_steps: The minimum number of decoding steps to take, i.e. the minimum length of
-        the predicted sequences. This does not include the start or end tokens. If `None`,
-        no minimum is enforced.
-
-    :param final_sequence_scorer: An optional `FinalSequenceScorer` which is used to score the final generated sequences.
-        The output from this module is what is returned by the `search` method. If not
-        specified, `SequenceLogProbabilityScorer` will be used, which scores the sequences
-        by the sum of the token log probabilities.
-
-    :param constraints: An optional list of `Constraint`s which should be applied during beam search. If not
-        provided, no constraints will be enforced.
-
-    """
-
-    def __init__(
-        self,
-        end_index: int,
-        *,
-        max_steps: int = 50,
-        beam_size: int = 10,
-        per_node_beam_size: Optional[int] = None,
-        sampler: Optional[Sampler] = None,
-        min_steps: Optional[int] = None,
-        final_sequence_scorer: Optional[FinalSequenceScorer] = None,
-        constraints: Optional[List[Constraint]] = None,
-    ) -> None:
-        if not max_steps > 0:
-            raise ValueError("max_steps must be positive")
-        if not beam_size > 0:
-            raise ValueError("beam_size must be positive")
-        if per_node_beam_size is not None and not per_node_beam_size > 0:
-            raise ValueError("per_node_beam_size must be positive")
-        if min_steps is not None:
-            if not min_steps >= 0:
-                raise ValueError("min_steps must be non-negative")
-            if not min_steps <= max_steps:
-                raise ValueError("min_steps must be less than or equal to max_steps")
-
-        self._end_index = end_index
-        self.max_steps = max_steps
-        self.beam_size = beam_size
-        self.per_node_beam_size = per_node_beam_size or beam_size
-        self.sampler = sampler or DeterministicSampler()
-        self.min_steps = min_steps or 0
-        self.final_sequence_scorer = final_sequence_scorer or SequenceLogProbabilityScorer()
-        self.constraints = constraints or []
-
-    @staticmethod
-    def _reconstruct_sequences(predictions, backpointers):
-        # Reconstruct the sequences.
-        # shape: [(batch_size, beam_size, 1)]
-        reconstructed_predictions = [predictions[-1].unsqueeze(2)]
-
-        if not backpointers:
-            return reconstructed_predictions
-
-        # shape: (batch_size, beam_size)
-        cur_backpointers = backpointers[-1]
-
-        for timestep in range(len(predictions) - 2, 0, -1):
-            # shape: (batch_size, beam_size, 1)
-            cur_preds = predictions[timestep].gather(1, cur_backpointers).unsqueeze(2)
-
-            reconstructed_predictions.append(cur_preds)
-
-            # shape: (batch_size, beam_size)
-            cur_backpointers = backpointers[timestep - 1].gather(1, cur_backpointers)
-
-        # shape: (batch_size, beam_size, 1)
-        final_preds = predictions[0].gather(1, cur_backpointers).unsqueeze(2)
-
-        reconstructed_predictions.append(final_preds)
-
-        return reconstructed_predictions
-
-    def search(
-        self,
-        start_predictions: torch.Tensor,
-        start_state: StateType,
-        step: StepFunctionType,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Given a starting state and a step function, apply beam search to find the
-        most likely target sequences.
-
-        Returns a tuple of `(predictions, final_scores)`, where `predictions`
-        has shape `(batch_size, beam_size, max_steps)` and `final_scores`
-        has shape `(batch_size, beam_size)`.
-
-        .. note::
-            If your step function returns `-inf` for some log probabilities
-            (like if you're using a masked log-softmax) then some of the "best"
-            sequences returned may also have `-inf` log probability. Specifically
-            this happens when the beam size is smaller than the number of actions
-            with finite log probability (non-zero probability) returned by the step function.
-            Therefore if you're using a mask you may want to check the results from `search`
-            and potentially discard sequences with non-finite log probability.
-
-        :param start_predictions: A tensor containing the initial predictions with shape `(batch_size,)`.
-            Usually the initial predictions are just the index of the "start" token
-            in the target vocabulary.
-
-        :param start_state: The initial state passed to the `step` function. Each value of the state dict
-            should be a tensor of shape `(batch_size, *)`, where `*` means any other
-            number of dimensions.
-
-        :param step: A function that is responsible for computing the next most likely tokens,
-            given the current state and the predictions from the last time step.
-            The function should accept two or three arguments:
-
-            - a tensor of shape `(group_size,)` or representing the index of the predicted
-            tokens from the last time step,
-            - the current state, a `StateType`, and
-            - optionally, the timestep, an `int`.
-
-            The `group_size` will be `batch_size * beam_size`, except in the initial
-            step, for which it will just be `batch_size`.
-
-            The function is expected to return a tuple, where the first element
-            is a tensor of shape `(group_size, vocab_size)` containing
-            the log probabilities of the tokens for the next step, and the second
-            element is the updated state. The tensor in the state should have shape
-            `(group_size, *)`, where `*` means any other number of dimensions.
-
-        """
-        step_signature = signature(step)
-        if len(step_signature.parameters) < 3:
-            # If the step function we're given does not take the time step argument, wrap it
-            # in one that does.
-            old_step = cast(StepFunctionTypeNoTimestep, step)
-
-            def new_step(last_predictions: torch.Tensor, state: Dict[str, torch.Tensor], time_step: int):
-                del time_step
-                return old_step(last_predictions, state)
-
-            return self._search(start_predictions, start_state, new_step)
-        else:
-            return self._search(start_predictions, start_state, cast(StepFunctionTypeWithTimestep, step))
-
-    def _search(
-        self,
-        start_predictions: torch.Tensor,
-        start_state: StateType,
-        step: StepFunctionTypeWithTimestep,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        batch_size = start_predictions.size()[0]
-
-        # List of (batch_size, beam_size) tensors. One for each time step. Does not
-        # include the start symbols, which are implicit.
-        predictions: List[torch.Tensor] = []
-
-        # List of (batch_size, beam_size) tensors. One for each time step. None for
-        # the first.  Stores the index n for the parent prediction, i.e.
-        # predictions[t-1][i][n], that it came from.
-        backpointers: List[torch.Tensor] = []
-
-        constraint_states = [constraint.init_state(batch_size) for constraint in self.constraints]
-
-        # Calculate the first timestep. This is done outside the main loop
-        # because we are going from a single decoder input (the output from the
-        # encoder) to the top `beam_size` decoder outputs. On the other hand,
-        # within the main loop we are going from the `beam_size` elements of the
-        # beam to `beam_size`^2 candidates from which we will select the top
-        # `beam_size` elements for the next iteration.
-        # shape: (batch_size, num_classes)
-        start_class_log_probabilities, state = step(start_predictions, start_state, 0)
-
-        num_classes = start_class_log_probabilities.size()[1]
-
-        # Make sure `per_node_beam_size` is not larger than `num_classes`.
-        if self.per_node_beam_size > num_classes:
-            raise ValueError(
-                f"Vocab size ({num_classes:d}) too small "
-                f"relative to per_node_beam_size ({self.per_node_beam_size:d}).\n"
-                f"Please decrease beam_size or per_node_beam_size."
-            )
-
-        sampler_state = self.sampler.init_state(start_class_log_probabilities, batch_size, num_classes)
-
-        # Apply all constraints.
-        if self.constraints:
-            # shape: (batch_size, 1, num_classes)
-            expanded_start_class_log_probabilities = start_class_log_probabilities.unsqueeze(1)
-            for constraint, constraint_state in zip(self.constraints, constraint_states):
-                expanded_start_class_log_probabilities = constraint.apply(
-                    constraint_state, expanded_start_class_log_probabilities
-                )
-            start_class_log_probabilities = expanded_start_class_log_probabilities.squeeze(1)
-
-        # Prevent selecting the end symbol if there is any min_steps constraint
-        if self.min_steps >= 1:
-            start_class_log_probabilities[:, self._end_index] = torch.finfo(
-                start_class_log_probabilities.dtype
-            ).min
-
-        # Get the initial predicted classed and their log probabilities.
-        # shape: (batch_size, beam_size), (batch_size, beam_size)
-        (
-            start_top_log_probabilities,
-            start_predicted_classes,
-            sampler_state,
-        ) = self.sampler.sample_beams(start_class_log_probabilities, self.beam_size, sampler_state)
-
-        if self.beam_size == 1 and (start_predicted_classes == self._end_index).all():
-            warnings.warn(
-                "Empty sequences predicted. You may want to increase the beam size or ensure "
-                "your step function is working properly.",
-                RuntimeWarning,
-            )
-            return start_predicted_classes.unsqueeze(-1), start_top_log_probabilities
-
-        # The log probabilities for the last time step.
-        # shape: (batch_size, beam_size)
-        last_log_probabilities = start_top_log_probabilities
-
-        # shape: [(batch_size, beam_size)]
-        predictions.append(start_predicted_classes)
-
-        # Log probability tensor that mandates that the end token is selected.
-        # shape: (batch_size * beam_size, num_classes)
-        log_probs_after_end = start_class_log_probabilities.new_full(
-            (batch_size * self.beam_size, num_classes),
-            torch.finfo(start_class_log_probabilities.dtype).min,
-        )
-        log_probs_after_end[:, self._end_index] = 0.0
-
-        # Set the same state for each element in the beam.
-        self._update_initial_state(state, batch_size)
-
-        for i, constraint in enumerate(self.constraints):
-            constraint_states[i] = constraint.update_state(constraint_states[i], start_predicted_classes)
-
-        for timestep in range(self.max_steps - 1):
-            # shape: (batch_size * beam_size,)
-            last_predictions = predictions[-1].reshape(batch_size * self.beam_size)
-
-            # If every predicted token from the last step is `self._end_index`,
-            # then we can stop early.
-            if (last_predictions == self._end_index).all():
-                break
-            # Take a step. This get the predicted log probs of the next classes
-            # and updates the state.
-            # shape: (batch_size * beam_size, num_classes)
-            class_log_probabilities, state = step(last_predictions, state, timestep + 1)
-
-            # Apply all constraints.
-            if self.constraints:
-                # shape: (batch_size, beam_size, num_classes)
-                reshaped_class_log_probabilities = class_log_probabilities.view(batch_size, self.beam_size, -1)
-                for constraint, constraint_state in zip(self.constraints, constraint_states):
-                    reshaped_class_log_probabilities = constraint.apply(
-                        constraint_state, reshaped_class_log_probabilities
-                    )
-                # shape: (batch_size * beam_size, num_classes)
-                class_log_probabilities = reshaped_class_log_probabilities.view(batch_size * self.beam_size, -1)
-
-            # The `timestep`-th iteration of the for loop is generating the `timestep + 2`-th token
-            # of the sequence (because `timestep` is 0-indexed and we generated the first token
-            # before the for loop). Here we block the end index if the search is not allowed to
-            # terminate on this iteration.
-            if timestep + 2 <= self.min_steps:
-                class_log_probabilities[:, self._end_index] = torch.finfo(class_log_probabilities.dtype).min
-
-            # shape: (batch_size * beam_size, num_classes)
-            last_predictions_expanded = last_predictions.unsqueeze(-1).expand(
-                batch_size * self.beam_size, num_classes
-            )
-
-            # Here we are finding any beams where we predicted the end token in
-            # the previous timestep and replacing the distribution with a
-            # one-hot distribution, forcing the beam to predict the end token
-            # this timestep as well.
-            # shape: (batch_size * beam_size, num_classes)
-            cleaned_log_probabilities = torch.where(
-                last_predictions_expanded == self._end_index,
-                log_probs_after_end,
-                class_log_probabilities,
-            )
-
-            # shape (both): (batch_size * beam_size, per_node_beam_size)
-            top_log_probabilities, predicted_classes, sampler_state = self.sampler.sample_nodes(
-                cleaned_log_probabilities, self.per_node_beam_size, sampler_state
-            )
-
-            # Here we expand the last log probabilities to (batch_size * beam_size, per_node_beam_size)
-            # so that we can add them to the current log probs for this timestep.
-            # This lets us maintain the log probability of each element on the beam.
-            # shape: (batch_size * beam_size, per_node_beam_size)
-            expanded_last_log_probabilities = (
-                last_log_probabilities.unsqueeze(2)
-                .expand(batch_size, self.beam_size, self.per_node_beam_size)
-                .reshape(batch_size * self.beam_size, self.per_node_beam_size)
-            )
-
-            # shape: (batch_size * beam_size, per_node_beam_size)
-            summed_top_log_probabilities = top_log_probabilities + expanded_last_log_probabilities
-
-            # shape: (batch_size, beam_size * per_node_beam_size)
-            reshaped_summed = summed_top_log_probabilities.reshape(
-                batch_size, self.beam_size * self.per_node_beam_size
-            )
-
-            # shape: (batch_size, beam_size * per_node_beam_size)
-            reshaped_predicted_classes = predicted_classes.reshape(
-                batch_size, self.beam_size * self.per_node_beam_size
-            )
-
-            # Keep only the top `beam_size` beam indices.
-            # shape (both): (batch_size, beam_size)
-            (
-                restricted_beam_log_probs,
-                restricted_beam_indices,
-                sampler_state,
-            ) = self.sampler.sample_beams(reshaped_summed, self.beam_size, sampler_state)
-
-            # Use the beam indices to extract the corresponding classes.
-            # shape: (batch_size, beam_size)
-            restricted_predicted_classes = reshaped_predicted_classes.gather(1, restricted_beam_indices)
-
-            predictions.append(restricted_predicted_classes)
-
-            # shape: (batch_size, beam_size)
-            last_log_probabilities = restricted_beam_log_probs
-
-            # The beam indices come from a `beam_size * per_node_beam_size` dimension where the
-            # indices with a common ancestor are grouped together. Hence
-            # dividing by per_node_beam_size gives the ancestor. (Note that this is integer
-            # division as the tensor is a LongTensor.)
-            # shape: (batch_size, beam_size)
-            backpointer = torch.divide(restricted_beam_indices, self.per_node_beam_size, rounding_mode="trunc")
-            backpointers.append(backpointer)
-
-            # Keep only the pieces of the state tensors corresponding to the
-            # ancestors created this iteration.
-            self._update_state(state, backpointer)
-
-            for i, constraint in enumerate(self.constraints):
-                constraint_states[i] = constraint.update_state(
-                    constraint_states[i], restricted_predicted_classes, last_backpointer=backpointer
-                )
-
-        # Warn about "-inf" log probabilities if not using any constraints (negligible
-        # log probabilities are expected when using constraints).
-        if not self.constraints and (
-            not torch.isfinite(last_log_probabilities).all()
-            or (last_log_probabilities == torch.finfo(last_log_probabilities.dtype).min).any()
-        ):
-            warnings.warn(
-                "Negligible log probabilities encountered ('-inf' or equivalent). "
-                "Some final sequences may not make sense. "
-                "This can happen when the beam size is larger than the number of valid (non-zero "
-                "probability) transitions that the step function produces.",
-                RuntimeWarning,
-            )
-
-        reconstructed_predictions = self._reconstruct_sequences(predictions, backpointers)
-
-        # shape: (batch_size, beam_size, max_steps)
-        all_predictions = torch.cat(list(reversed(reconstructed_predictions)), 2)
-
-        # Calculate the final sequence scores
-        # shape: (batch_size, beam_size)
-        final_scores = self.final_sequence_scorer.score(all_predictions, last_log_probabilities, self._end_index)
-
-        # Sort the sequences based on the final scores so the best scoring
-        # sequence is at index 0
-        sorted_final_scores, sorted_indices = torch.sort(final_scores, dim=1, descending=True)
-        sorted_all_predictions = torch.gather(
-            all_predictions, 1, sorted_indices.unsqueeze(-1).expand_as(all_predictions)
-        )
-
-        return sorted_all_predictions, sorted_final_scores
-
-    def _update_initial_state(self, state: StateType, batch_size: int):
-        """
-        Expand tensors in a state dictionary from `(batch_size, *)` to `(batch_size * beam_size, *)`.
-        """
-        for key, state_tensor in state.items():
-            if state_tensor is None:
-                continue
-            # shape: (batch_size * beam_size, *)
-            _, *last_dims = state_tensor.size()
-            state[key] = (
-                state_tensor.unsqueeze(1)
-                .expand(batch_size, self.beam_size, *last_dims)
-                .reshape(batch_size * self.beam_size, *last_dims)
-            )
-
-    def _update_state(self, state: StateType, backpointer: torch.Tensor):
-        batch_size = backpointer.size()[0]
-
-        for key, state_tensor in state.items():
-            if state_tensor is None:
-                continue
-            _, *last_dims = state_tensor.size()
-            # shape: (batch_size, beam_size, *)
-            expanded_backpointer = backpointer.view(batch_size, self.beam_size, *([1] * len(last_dims))).expand(
-                batch_size, self.beam_size, *last_dims
-            )
-            # shape: (batch_size * beam_size, *)
-            state[key] = (
-                state_tensor.reshape(batch_size, self.beam_size, *last_dims)
-                .gather(1, expanded_backpointer)
-                .reshape(batch_size * self.beam_size, *last_dims)
-            )

olmo_bitnet_1b/checkpoint.py

@@ -1,1671 +0,0 @@
-import gc
-import io
-import logging
-import pickle
-import shutil
-import traceback
-from abc import ABCMeta, abstractmethod
-from collections import defaultdict
-from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
-from contextlib import contextmanager
-from copy import deepcopy
-from dataclasses import dataclass, field, replace
-from functools import reduce
-from multiprocessing import shared_memory
-from pathlib import Path
-from typing import Any, Dict, Generator, List, Optional, Set, Tuple, cast
-
-import numpy as np
-import torch
-import torch.distributed.checkpoint as dist_cp
-import torch.multiprocessing as mp
-from packaging import version
-from torch.distributed import _remote_device
-from torch.distributed._shard._utils import narrow_tensor_by_index
-from torch.distributed._shard.metadata import ShardMetadata
-from torch.distributed._shard.sharded_tensor import ShardedTensor
-from torch.distributed.checkpoint.filesystem import WriteResult, _StorageInfo
-from torch.distributed.checkpoint.metadata import Metadata, MetadataIndex
-from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_dict
-from torch.distributed.checkpoint.planner import LoadItemType, ReadItem
-from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-from torch.distributed.fsdp import StateDictType
-from torch.distributed.fsdp.api import (
-    FullOptimStateDictConfig,
-    FullStateDictConfig,
-    ShardedOptimStateDictConfig,
-    ShardedStateDictConfig,
-)
-from torch.futures import Future
-
-try:
-    from torch.distributed.fsdp.flat_param import FlatParamHandle  # type: ignore
-except ModuleNotFoundError:
-    from torch.distributed.fsdp._flat_param import FlatParamHandle  # type: ignore
-
-from . import util
-
-from .aliases import PathOrStr
-from .config import BaseConfig, ShardedCheckpointerType, TrainConfig
-from .exceptions import OLMoCheckpointError
-from .optim import Optimizer, fix_optim_state_dict
-from .safetensors_util import safetensors_file_to_state_dict
-from .torch_util import (
-    barrier,
-    gc_cuda,
-    get_fs_local_rank,
-    get_global_rank,
-    get_world_size,
-)
-from .util import (
-    _get_s3_client,
-    default_thread_count,
-    dir_is_empty,
-    get_bytes_range,
-    get_progress_bar,
-    resource_path,
-    upload,
-    wait_for,
-)
-
-__all__ = [
-    "save_fsdp_model_and_optim_state",
-    "load_fsdp_model_and_optim_state",
-    "load_fsdp_optim_state",
-    "save_state_dict",
-    "load_state_dict",
-    "load_model_state",
-    "RemoteFileSystemWriter",
-    "RemoteFileSystemReader",
-    "Checkpointer",
-    "FullCheckpointer",
-    "TorchNewStyleShardedCheckpointer",
-    "TorchLegacyShardedCheckpointer",
-    "LocalShardedCheckpointer",
-    "build_sharded_checkpointer",
-]
-
-
-log = logging.getLogger(__name__)
-
-MODEL_AND_OPTIM_FOLDER = "model_and_optim"
-
-
-def save_fsdp_model_and_optim_state(
-    checkpoint_dir: PathOrStr,
-    fsdp_model: FSDP,
-    optim: Optimizer,
-    *,
-    upload_to: Optional[str] = None,
-    save_overwrite: bool = False,
-):
-    """
-    Use this to save a state dict for an FSDP model and its optimizer via :module:`torch.distributed.checkpoint`
-    functions. This should be used during distributed training and should be called by all ranks.
-
-    :param checkpoint_dir: The directory to save to.
-    :param fsdp_model: The FSDP model.
-    :param optim: The FSDP model's optimizer.
-    :param upload_to: Optional, a remote "directory" to upload the checkpoint files to.
-    :param save_overwrite: Overwrite existing files.
-
-    :raises FileExistsError: If a model and optim checkpoint already exists in ``checkpoint_dir`` and ``save_overwrite=False``.
-    """
-    checkpoint_dir = Path(checkpoint_dir)
-    target_dir = checkpoint_dir / MODEL_AND_OPTIM_FOLDER
-    if save_overwrite:
-        if get_fs_local_rank() == 0:
-            shutil.rmtree(target_dir, ignore_errors=True)
-    elif not dir_is_empty(target_dir):
-        raise FileExistsError(target_dir)
-    barrier()
-    if get_fs_local_rank() == 0:
-        target_dir.mkdir(exist_ok=True, parents=True)
-    barrier()
-    with FSDP.state_dict_type(
-        fsdp_model,
-        state_dict_type=StateDictType.SHARDED_STATE_DICT,
-        state_dict_config=ShardedStateDictConfig(offload_to_cpu=True),
-        optim_state_dict_config=ShardedOptimStateDictConfig(offload_to_cpu=True),
-    ):
-        model_and_optim_state = {
-            "model": fsdp_model.state_dict(),
-            "optim": FSDP.optim_state_dict(fsdp_model, optim),
-        }
-        dist_cp.save_state_dict(
-            model_and_optim_state,
-            RemoteFileSystemWriter(
-                target_dir,
-                upload_to=None if upload_to is None else f"{upload_to.rstrip('/')}/{MODEL_AND_OPTIM_FOLDER}",
-                save_overwrite=save_overwrite,
-            ),
-        )
-
-
-def load_fsdp_model_and_optim_state(
-    checkpoint_dir: PathOrStr,
-    fsdp_model: FSDP,
-    optim: Optimizer,
-    *,
-    local_cache: Optional[PathOrStr] = None,
-    load_optimizer_state: bool = True,
-):
-    """
-    Use this to load a state dict for an FSDP model and its optimizer via :module:`torch.distributed.checkpoint`
-    functions. This should be used during distributed training and should be called by all ranks.
-
-    :param checkpoint_dir: The checkpoint directory to load from. This can be a local or remote directory.
-    :param fsdp_model: The FSDP model.
-    :param optim: The FSDP model's optimizer.
-    :param local_cache: A local cache of the checkpoint directory. Use this when the ``checkpoint_dir`` is a
-        remote "directory" but there might be a cached version of the same artifacts.
-    :param load_optimizer_state: Set to ``False`` to skip loading the optimizer state.
-
-    :raises FileNotFoundError: If the ``checkpoint_dir`` doesn't contain a model and optimizer checkpoint.
-    """
-    load_path = str(checkpoint_dir).rstrip("/")
-    local_cache = None if local_cache is None else Path(local_cache)
-    with FSDP.state_dict_type(
-        fsdp_model,
-        state_dict_type=StateDictType.SHARDED_STATE_DICT,
-        state_dict_config=ShardedStateDictConfig(offload_to_cpu=True),
-        optim_state_dict_config=ShardedOptimStateDictConfig(offload_to_cpu=True),
-    ):
-        # Load the model state dict in place.
-        log.info("Loading model state...")
-        model_state = {"model": fsdp_model.state_dict()}
-        dist_cp.load_state_dict(
-            model_state,
-            RemoteFileSystemReader(
-                f"{load_path}/{MODEL_AND_OPTIM_FOLDER}",
-                local_cache=None if local_cache is None else local_cache / MODEL_AND_OPTIM_FOLDER,
-            ),
-        )
-        fsdp_model.load_state_dict(model_state["model"])
-
-        if not load_optimizer_state:
-            return
-
-        # Load optim state dict in place.
-        log.info("Loading sharded optimizer state...")
-        optim_state = load_sharded_optimizer_state_dict(
-            model_state_dict=model_state["model"],
-            optimizer_key="optim",
-            storage_reader=RemoteFileSystemReader(
-                f"{load_path}/{MODEL_AND_OPTIM_FOLDER}",
-                local_cache=None if local_cache is None else local_cache / MODEL_AND_OPTIM_FOLDER,
-            ),
-        )
-        del model_state
-        gc_cuda()
-        load_fsdp_optim_state(fsdp_model, optim, optim_state["optim"])
-
-
-def load_fsdp_optim_state(fsdp_model: FSDP, optim: Optimizer, optim_state: Dict[str, Any]):
-    log.info("Flattening sharded optimizer state...")
-    # NOTE: Careful! The order of the these arguments has changed from 2.0 to 2.1... ¯\_(ツ)_/¯
-    if version.parse(torch.__version__) < version.parse("2.1.0"):
-        flattened_osd = FSDP.optim_state_dict_to_load(optim_state, fsdp_model, optim)  # type: ignore
-    else:
-        flattened_osd = FSDP.optim_state_dict_to_load(fsdp_model, optim, optim_state)  # type: ignore
-    del optim_state
-    gc.collect()
-    log.info("Loading flattened optimizer state...")
-    # Put optim state on CPU since `Optimizer.load_state_dict()` will create a deepcopy of the whole state dict,
-    # which takes up unnecessary GPU memory.
-    for state in flattened_osd["state"].values():
-        for k in state.keys():
-            v = state[k]
-            if isinstance(v, torch.Tensor):
-                state[k] = v.to(device="cpu")
-    gc_cuda()
-    optim.load_state_dict(fix_optim_state_dict(optim, flattened_osd))
-
-
-def save_state_dict(
-    checkpoint_dir: PathOrStr,
-    fname: str,
-    state_dict: Dict[str, Any],
-    *,
-    upload_to: Optional[str] = None,
-    save_overwrite: bool = False,
-    synchronize: bool = True,
-):
-    """
-    Save a regular state dict to the file ``fname`` within ``checkpoint_dir`` using :func:`torch.save()`.
-    This can be used during distributed training or not. If during distributed training the ``fname`` should be unique
-    for each rank.
-
-    :param checkpoint_dir: The directory to save to.
-    :param fname: The target file within ``checkpoint_dir`` to save to. This should be a path relative to the ``checkpoint_dir``.
-    :param state_dict: The state dict to save.
-    :param upload_to: Optional, a remote "directory" to upload the file to.
-    :param save_overwrite: Overwrite existing files.
-    :param synchronize: If ``False``, don't do any distributed synchronization. Use this when only calling
-        this function from a single rank.
-
-    :raises FileExistsError: If the ``fname`` already exists within ``checkpoint_dir`` and ``save_overwrite=False``.
-    """
-    checkpoint_dir = Path(checkpoint_dir)
-    target_path = checkpoint_dir / fname
-    if save_overwrite:
-        target_path.unlink(missing_ok=True)
-    elif target_path.is_file():
-        raise FileExistsError(target_path)
-    if synchronize:
-        barrier()
-    target_path.parent.mkdir(exist_ok=True, parents=True)
-    if synchronize:
-        barrier()
-    torch.save(state_dict, target_path)
-    if upload_to is not None:
-        upload_target = f"{upload_to.rstrip('/')}/{fname}"
-        log.info(f"Uploading {target_path} to {upload_target}...")
-        upload(target_path, upload_target, save_overwrite=save_overwrite)
-
-
-def load_state_dict(
-    checkpoint_dir: PathOrStr,
-    fname: str,
-    *,
-    local_cache: Optional[PathOrStr] = None,
-    map_location: Optional[str] = None,
-):
-    """
-    Load a regular state dict from the file ``fname`` within ``checkpoint_dir`` using :func:`torch.load()`.
-    This can be used during distributed training or not.
-
-    :param checkpoint_dir: A local or remote checkpoint directory.
-    :param fname: The target file within the ``checkpoint_dir``. This should be a path relative to the ``checkpoint_dir``.
-    :param local_cache: A local cache of the checkpoint directory. Use this when the ``checkpoint_dir`` is a
-        remote "directory" but there might be a cached version of the same artifacts.
-
-    :raises FileNotFoundError: If ``fname`` doesn't exist in the ``checkpoint_dir`` or the local cache.
-    """
-    if fname.endswith(".pt"):
-        # Try safetensors version first.
-        try:
-            path = resource_path(
-                str(checkpoint_dir).rstrip("/"), fname[:-2] + "safetensors", local_cache=local_cache
-            )
-            return safetensors_file_to_state_dict(path, map_location=map_location)
-        except FileNotFoundError:
-            pass
-
-    path = resource_path(str(checkpoint_dir).rstrip("/"), fname, local_cache=local_cache)
-    return torch.load(path, map_location=map_location)
-
-
-def load_model_state(checkpoint_dir: PathOrStr, model: torch.nn.Module):
-    """
-    Load model state from a distributed FSDP model checkpoint created from :func:`save_fsdp_model_and_optim_state()`.
-    Note that ``model`` should not be wrapped with FSDP.
-    """
-    state_dict = {"model": model.state_dict()}
-    dist_cp.load_state_dict(
-        state_dict,
-        RemoteFileSystemReader(f"{str(checkpoint_dir).rstrip('/')}/{MODEL_AND_OPTIM_FOLDER}"),
-        no_dist=True,
-    )
-    model.load_state_dict(state_dict["model"])
-
-
-class RemoteFileSystemWriter(dist_cp.FileSystemWriter):
-    """
-    A subclass of :class:`~torch.distributed.checkpoint.FileSystemWriter` that can upload files
-    directly to a cloud bucket when ``upload_to`` is specified.
-    """
-
-    def __init__(
-        self,
-        path: PathOrStr,
-        single_file_per_rank: bool = True,
-        sync_files: bool = True,
-        thread_count: Optional[int] = None,
-        per_thread_copy_ahead: int = 10_000_000,
-        upload_to: Optional[str] = None,
-        save_overwrite: bool = False,
-    ) -> None:
-        if thread_count is not None and thread_count <= 0:
-            raise ValueError("thread count must be at least 1")
-        super().__init__(
-            path,
-            single_file_per_rank=single_file_per_rank,
-            sync_files=sync_files,
-            # NOTE: we default to 1 thread here instead of whatever `default_thread_count()`
-            # returns because uploading big checkpoint files with multiple threads causes
-            # boto3 to fail in weird ways.
-            thread_count=thread_count or 1,
-            per_thread_copy_ahead=per_thread_copy_ahead,
-        )
-        self.upload_to = None if upload_to is None else upload_to.rstrip("/")
-        self.save_overwrite = save_overwrite
-
-    def write_data(
-        self,
-        plan: dist_cp.SavePlan,
-        planner: dist_cp.SavePlanner,
-    ) -> Future[List[WriteResult]]:
-        fut = super().write_data(plan, planner)
-        if self.upload_to is not None:
-            files_to_upload = set()
-            for write_result in fut.wait():
-                files_to_upload.add(write_result.storage_data.relative_path)
-
-            # Create the global S3 client up front to work around a threading issue in boto.
-            if self.upload_to.startswith("s3://"):
-                _get_s3_client("s3")
-            elif self.upload_to.startswith("r2://"):
-                _get_s3_client("r2")
-
-            with ThreadPoolExecutor(max_workers=self.thread_count) as executor:
-                futures = []
-                for fname in files_to_upload:
-                    source = self.path / fname
-                    target = f"{self.upload_to}/{fname}"
-                    log.info(f"Uploading {source} to {target}...")
-                    futures.append(executor.submit(upload, source, target, save_overwrite=self.save_overwrite))
-                for f in as_completed(futures):
-                    try:
-                        f.result()
-                    except BaseException:
-                        # NOTE: we might get an error here that can't be pickled, which causes a different failure
-                        # later when PyTorch tries to reduce that error across ranks. So here we just make
-                        # sure we're raising a simple error type that can be pickled.
-                        raise OLMoCheckpointError(f"Original error:\n{traceback.format_exc()}")
-        return fut
-
-    def finish(self, metadata: Metadata, results: List[List[WriteResult]]) -> None:
-        super().finish(metadata, results)
-        if self.upload_to is not None:
-            source = self.path / ".metadata"
-            target = f"{self.upload_to}/.metadata"
-            log.info(f"Uploading {source} to {target}...")
-            upload(source, target, save_overwrite=self.save_overwrite)
-
-
-class RemoteFileSystemReader(dist_cp.StorageReader):
-    """
-    A :class:`~torch.distributed.checkpoint.StorageReader` based on :class:`~torch.distributed.checkpoint.FileSystemReader`
-    that can read data directly from cloud storage as well as a local directory.
-    """
-
-    def __init__(
-        self, path: PathOrStr, *, local_cache: Optional[PathOrStr] = None, thread_count: Optional[int] = None
-    ):
-        super().__init__()
-        if thread_count is not None and thread_count <= 0:
-            raise ValueError("thread count must be at least 1")
-        self.path = str(path).rstrip("/")
-        self.cache = None if local_cache is None else Path(local_cache)
-        self.thread_count = thread_count or default_thread_count()
-        self.storage_data: Dict[MetadataIndex, _StorageInfo] = dict()
-        self._metadata: Optional[Metadata] = None
-
-    def _get_bytes(self, relative_path: str, offset: int, length: int) -> bytes:
-        if self.cache is not None and (path := self.cache / relative_path).is_file():
-            return get_bytes_range(path, offset, length)
-        else:
-            return get_bytes_range(f"{self.path}/{relative_path}", offset, length)
-
-    def _get_content_for_read(self, read_item: ReadItem) -> Tuple[ReadItem, bytes]:
-        sinfo = self.storage_data[read_item.storage_index]
-        content = self._get_bytes(sinfo.relative_path, sinfo.offset, sinfo.length)
-        return (read_item, content)
-
-    def read_data(self, plan: dist_cp.LoadPlan, planner: dist_cp.LoadPlanner) -> Future[None]:
-        # Create the global S3 client up front to work around a threading issue in boto.
-        if isinstance(self.path, str):
-            if self.path.startswith("s3://"):
-                _get_s3_client("s3")
-            elif self.path.startswith("r2://"):
-                _get_s3_client("r2")
-
-        with ThreadPoolExecutor(max_workers=self.thread_count) as executor:
-            read_item_content_futures = []
-            for read_item in plan.items:
-                read_item_content_futures.append(executor.submit(self._get_content_for_read, read_item))
-            read_item_content_results = []
-            for f in as_completed(read_item_content_futures):
-                try:
-                    read_item_content_results.append(f.result())
-                except BaseException:
-                    # NOTE: we might get an error here that can't be pickled, which causes a different failure
-                    # later when PyTorch tries to reduce that error across ranks. So here we just make
-                    # sure we're raising a simple error type that can be pickled.
-                    raise OLMoCheckpointError(f"Original error:\n{traceback.format_exc()}")
-
-        # Modified from `FileSystemReader.read_data()`
-        for read_item, content in read_item_content_results:
-            bytes = io.BytesIO(content)
-            bytes.seek(0)
-            if read_item.type == LoadItemType.BYTE_IO:
-                planner.load_bytes(read_item, bytes)
-            else:
-                tensor = cast(torch.Tensor, torch.load(bytes, map_location="cpu"))
-                tensor = narrow_tensor_by_index(tensor, read_item.storage_offsets, read_item.lengths)
-                target_tensor = planner.resolve_tensor(read_item).detach()
-
-                assert (
-                    target_tensor.size() == tensor.size()
-                ), f"req {read_item.storage_index} mismatch sizes {target_tensor.size()} vs {tensor.size()}"
-                target_tensor.copy_(tensor)
-                planner.commit_tensor(read_item, target_tensor)
-
-        fut: Future = Future()
-        fut.set_result(None)
-        return fut
-
-    def read_metadata(self) -> Metadata:
-        if self._metadata is None:
-            with resource_path(self.path, ".metadata", local_cache=self.cache).open("rb") as metadata_file:
-                self._metadata = pickle.load(metadata_file)
-        return self._metadata
-
-    def set_up_storage_reader(self, metadata: Metadata, is_coordinator: bool) -> None:
-        del is_coordinator
-        self.storage_data = metadata.storage_data
-        assert self.storage_data is not None
-
-    def prepare_local_plan(self, plan: dist_cp.LoadPlan) -> dist_cp.LoadPlan:
-        return plan
-
-    def prepare_global_plan(self, global_plan: List[dist_cp.LoadPlan]) -> List[dist_cp.LoadPlan]:
-        return global_plan
-
-
-class Checkpointer(metaclass=ABCMeta):
-    def __init__(self, cfg: TrainConfig, thread_count: Optional[int] = None):
-        self.cfg = cfg
-        self.thread_count = thread_count or default_thread_count()
-
-    @abstractmethod
-    def save_checkpoint(
-        self,
-        dir: PathOrStr,
-        fsdp_model: FSDP,
-        optim: Optimizer,
-        train_state: Dict[str, Any],
-        *,
-        upload_to: Optional[str] = None,
-    ) -> None:
-        raise NotImplementedError
-
-    @abstractmethod
-    def restore_checkpoint(
-        self,
-        load_path: PathOrStr,
-        fsdp_model: FSDP,
-        optim: Optimizer,
-        *,
-        local_cache: Optional[PathOrStr] = None,
-        load_optimizer_state: bool = True,
-    ) -> Dict[str, Any]:
-        """
-        Restores a checkpoint to the model and optimizer. Returns the remaining trainer state.
-        """
-        raise NotImplementedError
-
-    def unshard_checkpoint(
-        self,
-        load_path: PathOrStr,
-        *,
-        local_cache: Optional[PathOrStr] = None,
-        load_optimizer_state: bool = True,
-        load_trainer_state: bool = True,
-        device: Optional[torch.device] = None,
-    ) -> Tuple[Dict[str, torch.Tensor], Optional[Dict[str, Any]], Optional[Dict[str, Any]]]:
-        """
-        Unshard a checkpoint.
-
-        Note this is not marked abstract because child classes are not required to implemented this.
-        """
-        del load_path, local_cache, load_optimizer_state, load_trainer_state, device
-        raise NotImplementedError
-
-    @contextmanager
-    def _temporary_wd(self, dir: PathOrStr) -> Generator[Path, None, None]:
-        # Make sure checkpoint directory doesn't exist unless it's okay to overwrite it.
-        checkpoint_dir = Path(dir)
-        if not dir_is_empty(checkpoint_dir):
-            if self.cfg.save_overwrite:
-                if get_fs_local_rank() == 0:
-                    shutil.rmtree(checkpoint_dir, ignore_errors=True)
-            else:
-                raise FileExistsError(checkpoint_dir)
-        # No need to mkdir here since we'll directly replace the temporary directory with
-        # this directory below.
-        barrier()
-
-        # Prepare temporary directory. We don't have to be as careful here, we can
-        # just remove it if it already exists.
-        checkpoint_dir_tmp = checkpoint_dir.with_name(checkpoint_dir.name + "-tmp")
-        if get_fs_local_rank() == 0:
-            shutil.rmtree(checkpoint_dir_tmp, ignore_errors=True)
-            checkpoint_dir_tmp.mkdir(exist_ok=True, parents=True)
-
-        barrier()
-
-        # Yield temporary directory for `.save_checkpoint()` to use.
-        yield checkpoint_dir_tmp
-
-        barrier()
-
-        # Finally if all went well replace the temporary directory with the actual
-        # checkpoint directory.
-        if get_fs_local_rank() == 0:
-            # Replace temp directory with target checkpoint directory.
-            try:
-                checkpoint_dir_tmp.replace(checkpoint_dir)
-            except FileNotFoundError:
-                # Caught when another (file-system) local rank 0 has already replaced the tmp directory.
-                # This can happen when nodes are saving to a common NFS drive but otherwise have distinct
-                # file-systems.
-                if not checkpoint_dir.exists():
-                    raise
-
-        # In the cases where we're using a shared NFS drive between ranks to save checkpoints,
-        # replacing the temp directory with the final directory from rank 0 might not be immediately
-        # realized in the file systems of the other ranks.
-        # So we wait here across all ranks until that final checkpoint directory is visible.
-        wait_for(lambda: checkpoint_dir.exists(), "Waiting for checkpoint directory", timeout=10.0)
-
-        barrier()
-
-    def _save_config(self, dir: PathOrStr, *, upload_to: Optional[str] = None) -> None:
-        if get_global_rank() == 0:
-            log.info("Saving config...")
-            self.cfg.save(config_path := Path(dir) / "config.yaml")
-            if upload_to is not None:
-                upload_target = f"{upload_to}/config.yaml"
-                log.info(f"Uploading {config_path} to {upload_target}")
-                upload(config_path, upload_target, save_overwrite=self.cfg.save_overwrite)
-
-
-class FullCheckpointer(Checkpointer):
-    """
-    A :class:`Checkpointer` that saves a single full model and optimizer state dictionary.
-    """
-
-    def save_checkpoint(
-        self,
-        dir: PathOrStr,
-        fsdp_model: FSDP,
-        optim: Optimizer,
-        trainer_state: Dict[str, Any],
-        *,
-        upload_to: Optional[str] = None,
-    ) -> None:
-        with self._temporary_wd(dir) as checkpoint_dir:
-            with FSDP.state_dict_type(
-                fsdp_model,
-                state_dict_type=StateDictType.FULL_STATE_DICT,
-                state_dict_config=FullStateDictConfig(rank0_only=True, offload_to_cpu=True),
-                optim_state_dict_config=FullOptimStateDictConfig(rank0_only=True, offload_to_cpu=True),
-            ):
-                # We'll write the model and optimizer state dicts individually to reduce (CPU) memory consumption.
-                # First the model state.
-                model_state_dict = fsdp_model.state_dict()
-                if get_global_rank() == 0:
-                    log.info("Saving model state...")
-                    save_state_dict(
-                        checkpoint_dir,
-                        "model.pt",
-                        model_state_dict,
-                        upload_to=upload_to,
-                        save_overwrite=self.cfg.save_overwrite,
-                        synchronize=False,
-                    )
-                del model_state_dict
-                barrier()
-
-                # Then the optimizer state.
-                optim_state_dict = FSDP.optim_state_dict(fsdp_model, optim)
-                if get_global_rank() == 0:
-                    log.info("Saving optim state...")
-                    save_state_dict(
-                        checkpoint_dir,
-                        "optim.pt",
-                        optim_state_dict,
-                        upload_to=upload_to,
-                        save_overwrite=self.cfg.save_overwrite,
-                        synchronize=False,
-                    )
-                del optim_state_dict
-                barrier()
-
-            # Save trainer state.
-            if get_global_rank() == 0:
-                log.info("Saving trainer state...")
-                save_state_dict(
-                    checkpoint_dir,
-                    "train.pt",
-                    trainer_state,
-                    upload_to=upload_to,
-                    save_overwrite=self.cfg.save_overwrite,
-                    synchronize=False,
-                )
-            # Save config.
-            self._save_config(checkpoint_dir, upload_to=upload_to)
-
-    def restore_checkpoint(
-        self,
-        load_path: PathOrStr,
-        fsdp_model: FSDP,
-        optim: Optimizer,
-        *,
-        local_cache: Optional[PathOrStr] = None,
-        load_optimizer_state: bool = True,
-    ) -> Dict[str, Any]:
-        with FSDP.state_dict_type(
-            fsdp_model,
-            state_dict_type=StateDictType.FULL_STATE_DICT,
-            state_dict_config=FullStateDictConfig(rank0_only=False, offload_to_cpu=True),
-            optim_state_dict_config=FullOptimStateDictConfig(rank0_only=False, offload_to_cpu=True),
-        ):
-            with torch.no_grad():
-                # fill everything with NaN, so we can check afterwards that every parameter has been restored
-                for module_name, module in fsdp_model.named_modules():
-                    if not isinstance(module, FSDP):
-                        continue
-                    for param in module.params:
-                        param.fill_(torch.nan)
-
-                # restore params from checkpoint
-                state_dict_to_load = load_state_dict(
-                    load_path, "model.pt", local_cache=local_cache, map_location="cpu"
-                )
-                (
-                    state_dict_to_load,
-                    og_keys_to_new,
-                ) = fsdp_model._fsdp_wrapped_module._make_state_dict_compatible(state_dict_to_load)
-
-                for module_name, module in fsdp_model.named_modules():
-                    if not isinstance(module, FSDP):
-                        continue
-                    for param in module.params:
-                        assert param._is_flat_param
-                        for fqn, spi in zip(param._fqns, param._shard_param_infos):
-                            if not spi.in_shard:
-                                continue
-                            key = f"{module_name}.{fqn}"
-                            key = key.replace("_fsdp_wrapped_module.", "")
-                            key = key.lstrip(".")
-                            t = state_dict_to_load[key]
-                            t = t.flatten()
-                            param[spi.offset_in_shard : spi.offset_in_shard + spi.numel_in_shard].copy_(
-                                t[spi.intra_param_start_idx : spi.intra_param_end_idx + 1]
-                            )
-
-                # make sure that every parameter has been restored
-                for module_name, module in fsdp_model.named_modules():
-                    if not isinstance(module, FSDP):
-                        continue
-                    for param in module.params:
-                        if torch.isnan(param).any():
-                            raise ValueError(
-                                f"Module '{module_name}' contains NaNs, this is likely a bug restoring from full checkpoints"
-                            )
-
-            # Load optimizer state.
-            if load_optimizer_state:
-                optim_state_dict_to_load = load_state_dict(
-                    load_path, "optim.pt", local_cache=local_cache, map_location="cpu"
-                )
-                optim_state_dict_to_load = self._make_optim_state_dict_compatible(
-                    optim_state_dict_to_load,
-                    og_keys_to_new,
-                )
-                load_fsdp_optim_state(fsdp_model, optim, optim_state_dict_to_load)
-                del optim_state_dict_to_load
-
-            # Load other state.
-            try:
-                trainer_state = load_state_dict(load_path, "train.pt", local_cache=local_cache)
-            except FileNotFoundError:
-                # for backwards compatibility
-                trainer_state = load_state_dict(load_path, "other.pt", local_cache=local_cache)
-        barrier()
-        return trainer_state
-
-    def _make_optim_state_dict_compatible(
-        self, optim_state_dict: Dict[str, Any], og_keys_to_new: Dict[str, Set[str]]
-    ) -> Dict[str, Any]:
-        # This state dict comes in two forms: one where the state keys are integers and one where the
-        # keys are fully qualified parameter names. The latter case is easier to deal with here so we
-        # first transform the integer key form into the FQN key form.
-        if isinstance(optim_state_dict["param_groups"][0]["params"][0], int):
-            id_to_fqn: Dict[int, str] = {}
-            for group in optim_state_dict["param_groups"]:
-                new_param_names = []
-                for fqn, id in zip(group["param_names"], group["params"]):
-                    fqn = fqn.replace("_fsdp_wrapped_module.", "")
-                    id_to_fqn[id] = fqn
-                    new_param_names.append(fqn)
-                group["param_names"] = new_param_names
-                group["params"] = new_param_names
-            for id in list(optim_state_dict["state"].keys()):
-                optim_state_dict["state"][id_to_fqn[id]] = optim_state_dict["state"].pop(id)
-        else:
-            # Otherwise we still want to clean up the param names to remove the "_fsdp_wrapped_module." prefix.
-            for group in optim_state_dict["param_groups"]:
-                group["param_names"] = [fqn.replace("_fsdp_wrapped_module.", "") for fqn in group["param_names"]]
-                group["params"] = [fqn.replace("_fsdp_wrapped_module.", "") for fqn in group["params"]]
-                assert group["param_names"] == group["params"]
-            for key in list(optim_state_dict["state"].keys()):
-                optim_state_dict["state"][key.replace("_fsdp_wrapped_module.", "")] = optim_state_dict[
-                    "state"
-                ].pop(key)
-
-        # Now we can transform the state dict by renaming parameters according to `og_keys_to_new`.
-        # First fix param names in the state.
-        for og_key, new_keys in og_keys_to_new.items():
-            og_state = optim_state_dict["state"].pop(og_key, None)
-            if og_state is None:
-                continue
-            for i, new_key in enumerate(new_keys):
-                if i == len(new_keys) - 1:
-                    optim_state_dict["state"][new_key] = og_state
-                else:
-                    optim_state_dict["state"][new_key] = deepcopy(og_state)
-        # Now fix param names in the param groups.
-        for group in optim_state_dict["param_groups"]:
-            og_names = group["params"]
-            new_names = []
-            for og_key in og_names:
-                for new_key in og_keys_to_new[og_key]:
-                    new_names.append(new_key)
-            group["params"] = new_names
-            group["param_names"] = new_names
-
-        return optim_state_dict
-
-    def load_checkpoint(
-        self,
-        load_path: PathOrStr,
-        *,
-        local_cache: Optional[PathOrStr] = None,
-        load_optimizer_state: bool = True,
-        device: Optional[torch.device] = None,
-    ) -> Tuple[Dict[str, torch.Tensor], Optional[Dict[str, Any]]]:
-        device = device if device is not None else torch.device("cpu")
-        model_state = load_state_dict(load_path, "model.pt", local_cache=local_cache, map_location=device)  # type: ignore
-        optim_state = None
-        if load_optimizer_state:
-            optim_state = load_state_dict(load_path, "optim.pt", local_cache=local_cache, map_location=device)  # type: ignore
-        return model_state, optim_state
-
-
-class TorchNewStyleShardedCheckpointer(Checkpointer):
-    """
-    A sharded :class:`Checkpointer` that uses PyTorch's new distributed checkpointing functionality.
-    """
-
-    def save_checkpoint(
-        self,
-        dir: PathOrStr,
-        fsdp_model: FSDP,
-        optim: Optimizer,
-        trainer_state: Dict[str, Any],
-        *,
-        upload_to: Optional[str] = None,
-    ) -> None:
-        with self._temporary_wd(dir) as checkpoint_dir:
-            # Save model and optim state.
-            save_fsdp_model_and_optim_state(
-                checkpoint_dir,
-                fsdp_model,
-                optim,
-                upload_to=upload_to,
-                save_overwrite=self.cfg.save_overwrite,
-            )
-
-            # Save trainer state.
-            log.info("Saving trainer state...")
-            save_state_dict(
-                checkpoint_dir,
-                f"train/rank{get_global_rank()}.pt",
-                trainer_state,
-                upload_to=upload_to,
-                save_overwrite=self.cfg.save_overwrite,
-            )
-
-            # Save config.
-            self._save_config(checkpoint_dir, upload_to=upload_to)
-
-    def restore_checkpoint(
-        self,
-        load_path: PathOrStr,
-        fsdp_model: FSDP,
-        optim: Optimizer,
-        *,
-        local_cache: Optional[PathOrStr] = None,
-        load_optimizer_state: bool = True,
-    ) -> Dict[str, Any]:
-        # Load model and optimizer state in place.
-        log.info("Loading model and optimizer state...")
-        load_fsdp_model_and_optim_state(
-            load_path,
-            fsdp_model,
-            optim,
-            local_cache=local_cache,
-            load_optimizer_state=load_optimizer_state,
-        )
-
-        # Load trainer state dict.
-        log.info("Loading trainer state...")
-        try:
-            trainer_state = load_state_dict(
-                load_path, f"train/rank{get_global_rank()}.pt", local_cache=local_cache
-            )
-        except FileNotFoundError:
-            # Fall back to rank 0 train state.
-            # This can happen when we're restoring a checkpoint with a different world size.
-            trainer_state = load_state_dict(load_path, "train/rank0.pt", local_cache=local_cache)
-        barrier()
-        return trainer_state
-
-
-class TorchLegacyShardedCheckpointer(Checkpointer):
-    """
-    A sharded :class:`Checkpointer` that just uses `torch.save()` with extra logic for handling FSDP model
-    and optim state.
-
-    The world size must be kept consistent when using this checkpointer.
-    """
-
-    def save_checkpoint(
-        self,
-        dir: PathOrStr,
-        fsdp_model: FSDP,
-        optim: Optimizer,
-        trainer_state: Dict[str, Any],
-        *,
-        upload_to: Optional[str] = None,
-    ) -> None:
-        with self._temporary_wd(dir) as checkpoint_dir:
-            with FSDP.state_dict_type(
-                fsdp_model,
-                state_dict_type=StateDictType.SHARDED_STATE_DICT,
-                state_dict_config=ShardedStateDictConfig(offload_to_cpu=True),
-                optim_state_dict_config=ShardedOptimStateDictConfig(offload_to_cpu=True),
-            ):
-                state_dict = {
-                    "model": fsdp_model.state_dict(),
-                    "optim": FSDP.optim_state_dict(fsdp_model, optim),
-                    **trainer_state,
-                }
-                save_state_dict(
-                    checkpoint_dir,
-                    f"rank{get_global_rank()}.pt",
-                    state_dict,
-                    upload_to=upload_to,
-                    save_overwrite=self.cfg.save_overwrite,
-                )
-
-            # Save config.
-            self._save_config(checkpoint_dir, upload_to=upload_to)
-
-    def restore_checkpoint(
-        self,
-        load_path: PathOrStr,
-        fsdp_model: FSDP,
-        optim: Optimizer,
-        *,
-        local_cache: Optional[PathOrStr] = None,
-        load_optimizer_state: bool = True,
-    ) -> Dict[str, Any]:
-        with FSDP.state_dict_type(
-            fsdp_model,
-            state_dict_type=StateDictType.SHARDED_STATE_DICT,
-            state_dict_config=ShardedStateDictConfig(offload_to_cpu=True),
-            optim_state_dict_config=ShardedOptimStateDictConfig(offload_to_cpu=True),
-        ):
-            # Deserialize state dict.
-            state_dict = load_state_dict(
-                load_path, f"rank{get_global_rank()}.pt", local_cache=local_cache, map_location="cpu"
-            )
-
-            # Load model and optimizer state.
-            log.info("Loading model state...")
-            fsdp_model.load_state_dict(state_dict["model"])
-            del state_dict["model"]
-            if load_optimizer_state:
-                log.info("Loading optimizer state...")
-                load_fsdp_optim_state(fsdp_model, optim, state_dict["optim"])
-            del state_dict["optim"]
-
-        barrier()
-        return state_dict
-
-    def unshard_checkpoint(
-        self,
-        load_path: PathOrStr,
-        *,
-        local_cache: Optional[PathOrStr] = None,
-        load_optimizer_state: bool = True,
-        load_trainer_state: bool = True,
-        device: Optional[torch.device] = None,
-    ) -> Tuple[Dict[str, torch.Tensor], Optional[Dict[str, Any]], Optional[Dict[str, Any]]]:
-        assert local_cache is None, "this method currently only supports local files"
-        full_state_dict = self._unshard(load_path, device or torch.device("cpu"), skip_keys={"rng"})
-        model_state = full_state_dict.pop("model")
-        optim_state = full_state_dict.pop("optim")
-        return (
-            model_state,
-            optim_state if load_optimizer_state else None,
-            full_state_dict if load_trainer_state else None,
-        )
-
-    def _copy_sharded_tensors_to_shared_mem(self, state: Dict, world_size: int, rank: int, key: Tuple):
-        key = tuple() if key is None else key
-        if isinstance(state, (list, tuple, set)):
-            for i, sub_state in enumerate(state):
-                self._copy_sharded_tensors_to_shared_mem(sub_state, world_size, rank, key + (i,))
-        elif isinstance(state, dict):
-            for name in state.keys():
-                self._copy_sharded_tensors_to_shared_mem(state[name], world_size, rank, key + (name,))
-        elif isinstance(state, ShardedTensor):
-            self._copy_sharded_tensor_to_shared_mem(state, world_size, rank, key)
-            return
-        else:
-            return
-
-    def _get_shard_placement_and_rank_sizes(
-        self, shards_metadata: List[ShardMetadata], world_size: int
-    ) -> Tuple[Dict[ShardMetadata, Tuple[int, int]], List[int]]:
-        def shard_size(shard_md):
-            return reduce((lambda x, y: x * y), shard_md.shard_sizes)  # type: ignore[attr-defined]
-
-        rank_sizes = [0 for _ in range(world_size)]
-        shard_placement: Dict[ShardMetadata, Tuple[int, int]] = {}
-        for shard_md in shards_metadata:
-            shard_rank = cast(_remote_device, shard_md.placement).rank()
-            assert shard_rank is not None
-            if shard_rank >= world_size:
-                raise RuntimeError(f"Shard rank {shard_rank} exceeds world size {world_size}")
-
-            shard_placement[shard_md] = (shard_rank, rank_sizes[shard_rank])
-            rank_sizes[shard_rank] += shard_size(shard_md)
-
-        return shard_placement, rank_sizes
-
-    def _copy_sharded_tensor_to_shared_mem(
-        self, sharded_tensor: ShardedTensor, world_size: int, rank: int, key: Tuple
-    ) -> Any:
-        shard0_md = sharded_tensor.metadata()
-        shard_placement, rank_sizes = self._get_shard_placement_and_rank_sizes(
-            shard0_md.shards_metadata, world_size
-        )
-
-        rank_size = rank_sizes[rank]
-        assert rank_size >= 0
-        if rank_size == 0:
-            return
-
-        assert shard0_md.tensor_properties.dtype == torch.float32, "Expected sharded tensor to be fp32"
-        numpy_type = np.float32
-
-        sharded_memory_name = "-".join(key + (str(rank),))
-
-        shm = shared_memory.SharedMemory(
-            create=True, size=rank_size * np.dtype(numpy_type).itemsize, name=sharded_memory_name
-        )
-        np_arr = np.ndarray((rank_size,), dtype=numpy_type, buffer=shm.buf)
-
-        for local_shard in sharded_tensor.local_shards():
-            shard_rank = cast(_remote_device, local_shard.metadata.placement).rank()
-            assert shard_rank == rank
-
-            src = local_shard.tensor.flatten()
-            shard_offset = shard_placement[local_shard.metadata][1]
-
-            np_arr[shard_offset : shard_offset + src.numel()] = src.numpy()
-
-        shm.close()
-
-    def _copy_sharded_data_to_shared_mem(self, world_size: int, shard_filepath: Path):
-        shard_number = int(shard_filepath.name[4:-3])
-        log.info("Starting unsharding shard number %d to shared memory", shard_number)
-
-        with self._patch_sharded_tensor_load():
-            shard = torch.load(shard_filepath, map_location="cpu")
-            log.debug("Done loading shard number %d", shard_number)
-
-        self._copy_sharded_tensors_to_shared_mem(
-            shard, world_size, shard_number, (str(shard_filepath.parent).replace("/", "_"),)
-        )
-        log.info("Done unsharding shard number %d to shared memory", shard_number)
-
-    def _unshard_using_sharded_mem(
-        self, state: Any, world_size: int, device: torch.device, shard_dir: PathOrStr
-    ) -> Any:
-        return self._unshard_state_using_shared_mem(state, world_size, device, (str(shard_dir).replace("/", "_"),))
-
-    def _unshard_state_using_shared_mem(
-        self, state: Any, world_size: int, device: torch.device, key: Tuple
-    ) -> Any:
-        if isinstance(state, (list, tuple, set)):
-            return state.__class__(
-                self._unshard_state_using_shared_mem(sub_state, world_size, device, key + (i,))
-                for i, sub_state in enumerate(state)
-            )
-        elif isinstance(state, dict):
-            return {
-                name: self._unshard_state_using_shared_mem(state[name], world_size, device, key + (name,))
-                for name in state.keys()
-            }
-        elif isinstance(state, ShardedTensor):
-            return self._unshard_tensor_using_shared_mem(state, world_size, device, key)
-        elif isinstance(state, torch.Tensor):
-            return state.to(device=device)
-        else:
-            return state
-
-    def _unshard_tensor_using_shared_mem(
-        self, sharded_tensor: ShardedTensor, world_size: int, device: torch.device, key: Tuple
-    ) -> torch.Tensor:
-        shard0_md = sharded_tensor.metadata()
-
-        def shard_size(shard_md):
-            return reduce((lambda x, y: x * y), shard_md.shard_sizes)  # type: ignore[attr-defined]
-
-        shard_placement, rank_sizes = self._get_shard_placement_and_rank_sizes(
-            shard0_md.shards_metadata, world_size
-        )
-
-        assert shard0_md.tensor_properties.dtype == torch.float32, "Expected sharded tensor to be fp32"
-        numpy_type = np.float32
-
-        out = torch.empty(
-            *sharded_tensor.metadata().size, dtype=sharded_tensor.metadata().tensor_properties.dtype, device=device
-        )
-        dims = len(sharded_tensor.metadata().size)
-        for shard_md, (rank, rank_offset) in shard_placement.items():
-            if rank >= world_size:
-                raise RuntimeError(f"Shard rank {rank} exceeds world size {world_size}")
-
-            sharded_memory_name = "-".join(key + (str(rank),))
-            shm = shared_memory.SharedMemory(name=sharded_memory_name)
-
-            rank_size = rank_sizes[rank]
-            assert rank_size >= 0
-            if rank_size == 0:
-                continue
-
-            np_arr = np.ndarray((rank_size,), dtype=numpy_type, buffer=shm.buf)
-
-            tensor = torch.from_numpy(np_arr)[rank_offset : rank_offset + shard_size(shard_md)]
-            tensor = tensor.view(shard_md.shard_sizes)
-
-            out_narrow_view = out
-            for dim in range(dims):
-                out_narrow_view = out_narrow_view.narrow(
-                    dim,
-                    shard_md.shard_offsets[dim],
-                    shard_md.shard_sizes[dim],
-                )
-
-            out_narrow_view.copy_(tensor)
-
-            shm.close()
-            shm.unlink()
-
-        return out
-
-    @contextmanager
-    def _patch_sharded_tensor_load(self):
-        """
-        Monkeypatch for torch's ShardedTensor, so we can unpickle without having torch.distributed set up.
-        """
-
-        def _rebuild_from_type_v2_monkey(func, new_type, args, state):
-            ret = func(*args)
-            if type(ret) is not new_type:
-                ret = ret.as_subclass(new_type)
-
-            # Shortcut the construction of ShardedTensor
-            # This is in the top 5 of my worst hacks.
-            if isinstance(ret, ShardedTensor):
-                ret._local_shards, ret._metadata, _, ret._sharding_spec, ret._init_rrefs = state
-                return ret
-
-            # The rest of this function ought to be in the top 5 of somebody else's worst hacks.
-            # Tensor does define __setstate__ even though it doesn't define
-            # __getstate__. So only use __setstate__ if it is NOT the one defined
-            # on Tensor
-            if getattr(ret.__class__, "__setstate__", torch.Tensor.__setstate__) is not torch.Tensor.__setstate__:
-                ret.__setstate__(state)
-            else:
-                ret = torch._utils._set_obj_state(ret, state)
-            return ret
-
-        original_rebuild_from_type_v2 = torch._tensor._rebuild_from_type_v2
-        try:
-            torch._tensor._rebuild_from_type_v2 = _rebuild_from_type_v2_monkey
-            yield
-        finally:
-            torch._tensor._rebuild_from_type_v2 = original_rebuild_from_type_v2
-
-    def _unshard(self, input_dir: PathOrStr, device: torch.device, skip_keys: Optional[Set[str]] = None):
-        """
-        The current unsharding implementation consists of:
-
-        1. Loading each shard on a separate process and copying their sharded tensors to shared memory.
-        2. Loading 1 shard on the main process as a base unsharded object.
-        3. Using the sharded tensors in shared memory to populate the base unsharded object.
-
-        This implementation replaced a prior implementation that instead loaded
-        all shards using threads, because that implementation turned out to
-        be extremely slow (e.g. 6+ hours) sometimes when the world size was 1024.
-        The current implementation is slower than the old one in many scenarios,
-        but is significantly faster in the above mentioned case (e.g. 30 minutes)
-        if there are enough CPUs.
-        """
-
-        input_dir = Path(input_dir)
-        skip_keys = skip_keys or set()
-
-        shard_filepaths = list(input_dir.glob("rank*.pt"))
-        world_size = len(shard_filepaths)
-        if world_size == 0:
-            raise RuntimeError("No shards found for unsharding")
-
-        log.info("Number of shards: %d", world_size)
-        shard_size_gb = shard_filepaths[0].stat().st_size / (1024 * 1024 * 1024)
-        min_ram_required_estimate_gb = shard_size_gb * world_size
-        log.info(
-            "Shards are %.2fGB each, at least %.2fGB RAM is required", shard_size_gb, min_ram_required_estimate_gb
-        )
-
-        log.info("Copying sharded tensors to shared memory using multiple processes")
-        # Copy sharded data to shared memory using multiple processes, so this process can load
-        # from memory rather than disk. We spawn a new process instead of forking since shared memory
-        # appears to get deleted when forked processes end for some reason.
-        executor = ProcessPoolExecutor(
-            mp_context=mp.get_context("spawn"), initializer=util.prepare_cli_environment
-        )
-        futures = []
-        for shard_filepath in shard_filepaths:
-            shard_rank = int(shard_filepath.name[4:-3])
-
-            if shard_rank >= world_size:
-                raise RuntimeError(
-                    f"Shard rank {shard_rank} of file {shard_filepath} exceeds world size {world_size}"
-                )
-
-            futures.append(executor.submit(self._copy_sharded_data_to_shared_mem, world_size, shard_filepath))
-
-        for f in as_completed(futures):
-            f.result()
-        executor.shutdown()
-
-        log.info("Loading a shard on the main process to be unsharded state")
-        with self._patch_sharded_tensor_load():
-            state = torch.load(shard_filepaths[0], map_location="cpu")
-
-        for key in skip_keys:
-            if key in state:
-                del state[key]
-
-        log.info("Unsharding from %d shards ...", world_size)
-        return self._unshard_using_sharded_mem(state, world_size, device, input_dir)
-
-
-@dataclass
-class _LocalShardedCheckpointerMetadata(BaseConfig):
-    world_size: int = field(default_factory=get_world_size)
-
-
-@dataclass
-class _FlatParamShard:
-    full_shape: torch.Size
-    shard_offsets: Tuple[int, int]
-    shard_data: Optional[torch.Tensor]
-
-    def copy_into(self, full_tensor: torch.Tensor) -> None:
-        assert self.shard_data is not None
-        full_tensor_shard_view = full_tensor.view(-1)[self.shard_offsets[0] : self.shard_offsets[1] + 1]
-        assert self.shard_data.shape == full_tensor_shard_view.shape
-        full_tensor_shard_view.copy_(self.shard_data)
-
-
-class LocalShardedCheckpointer(Checkpointer):
-    """
-    A sharded :class:`Checkpointer` that directly saves the local FSDP flat params data.
-    The optimizer state is saved directly with `torch.save()` without reformatting via FSDP methods.
-
-    The world size must be kept consistent when using this checkpointer. However, you can easily
-    reconstruct a full unsharded model and/or optimizer state dictionary from a single Python process
-    using :meth:`unshard_checkpoint()` (no distributed initialization required).
-    """
-
-    # These correspond to metadata attributes on `torch.distributed.fsdp.flat_param.FlatParameter`.
-    _FLAT_PARAM_METADATA_TO_SAVE = (
-        "_fqns",
-        "_shard_param_offsets",
-        "_shard_indices",
-        "_numels",
-        "_numels_with_padding",
-        "_shapes",
-        "_shard_numel_padded",
-        "_shard_param_infos",
-    )
-
-    def _fsdp_modules(self, fsdp_model: FSDP) -> List[Tuple[str, FSDP]]:
-        """
-        Returns a list of FSDP modules with their FQN.
-        """
-        modules = []
-        for name, module in fsdp_model.named_modules():
-            if isinstance(module, FSDP):
-                modules.append((name, module))
-        return modules
-
-    def _prepare_fsdp_model(self, fsdp_model: FSDP) -> None:
-        from torch.distributed.fsdp._runtime_utils import _lazy_init
-
-        # TODO (epwalsh): I'm not sure if this is necessary, but this is what PyTorch does before saving/loading
-        # an FSDP state dict through the built-in methods.
-        if torch.cuda.is_available():
-            torch.cuda.synchronize()
-        _lazy_init(fsdp_model, fsdp_model)
-
-    def _fsdp_handles(self, fsdp_model: FSDP) -> List[FlatParamHandle]:
-        if version.parse(torch.__version__) < version.parse("2.1.0"):
-            return fsdp_model._handles  # type: ignore
-        elif version.parse(torch.__version__) < version.parse("2.3.0"):
-            # Handle could be None if the FSDP wrapper doesn't manage any parameters.
-            if hasattr(fsdp_model, "_handle") and fsdp_model._handle is not None:
-                return [fsdp_model._handle]  # type: ignore
-            else:
-                return []
-        else:
-            # Need to verify FSDP internals with newer versions.
-            raise NotImplementedError
-
-    @torch.no_grad()
-    def _get_flat_param_state_to_save(self, fsdp_model: FSDP) -> Dict[str, Any]:
-        self._prepare_fsdp_model(fsdp_model)
-        module_data = []
-        for module_fqn, fsdp_module in self._fsdp_modules(fsdp_model):
-            handle_data = []
-            for handle in self._fsdp_handles(fsdp_module):
-                data: Dict[str, Any] = {}
-                # This is a `FlatParameter` instance.
-                # See `torch.distributed.fsdp.flat_param` for the API.
-                flat_param = handle.flat_param
-                data["flat_param.data"] = flat_param.detach()
-                for key in self._FLAT_PARAM_METADATA_TO_SAVE:
-                    if hasattr(flat_param, key):
-                        data[f"flat_param.{key}"] = getattr(flat_param, key)
-                handle_data.append(data)
-            module_data.append({"handles": handle_data, "name": module_fqn})
-        return {"modules": module_data}
-
-    @torch.no_grad()
-    def _load_flat_param_state(self, fsdp_model: FSDP, model_state: Dict[str, Any]):
-        """Load the state produced from `self._get_flat_param_state_to_save()`."""
-        self._prepare_fsdp_model(fsdp_model)
-        fsdp_modules = self._fsdp_modules(fsdp_model)
-        assert len(model_state["modules"]) == len(fsdp_modules)
-        for (_, fsdp_module), module_data in zip(fsdp_modules, model_state["modules"]):
-            handles = self._fsdp_handles(fsdp_module)
-            assert len(handles) == len(module_data["handles"])
-            for handle, data in zip(handles, module_data["handles"]):
-                flat_param = handle.flat_param
-                # Make sure metadata matches.
-                for key in self._FLAT_PARAM_METADATA_TO_SAVE:
-                    if hasattr(flat_param, key):
-                        assert getattr(flat_param, key) == data[f"flat_param.{key}"]
-                # Load the flat sharded data.
-                flat_param.copy_(data["flat_param.data"])
-
-    def _save_metadata(self, dir: PathOrStr, *, upload_to: Optional[str] = None) -> None:
-        if get_fs_local_rank() == 0:
-            log.info("Saving metadata...")
-            metadata = _LocalShardedCheckpointerMetadata()
-            metadata.save(metadata_path := Path(dir) / "metadata.yaml")
-            if upload_to is not None and get_global_rank() == 0:
-                upload_target = f"{upload_to}/metadata.yaml"
-                log.info(f"Uploading {metadata_path} to {upload_target}")
-                upload(metadata_path, upload_target, save_overwrite=self.cfg.save_overwrite)
-
-    def _load_metadata(
-        self, load_path: PathOrStr, *, local_cache: Optional[PathOrStr] = None
-    ) -> _LocalShardedCheckpointerMetadata:
-        metadata_path = resource_path(load_path, "metadata.yaml", local_cache=local_cache)
-        return _LocalShardedCheckpointerMetadata.load(metadata_path)
-
-    def save_checkpoint(
-        self,
-        dir: PathOrStr,
-        fsdp_model: FSDP,
-        optim: Optimizer,
-        trainer_state: Dict[str, Any],
-        *,
-        upload_to: Optional[str] = None,
-    ) -> None:
-        with self._temporary_wd(dir) as checkpoint_dir:
-            # Gather local FSDP flat params data to save.
-            # We also save some flat param metadata like the corresponding fully qualified names (fqns)
-            # of each original parameter so we can validate that the sharding is the same when loading
-            # one of these checkpoints.
-            log.info("Saving local FSDP flat params data...")
-            save_state_dict(
-                checkpoint_dir,
-                f"model/rank{get_global_rank()}.pt",
-                self._get_flat_param_state_to_save(fsdp_model),
-                upload_to=upload_to,
-                save_overwrite=self.cfg.save_overwrite,
-            )
-
-            # Save optimizer state.
-            log.info("Saving local optimizer state...")
-            save_state_dict(
-                checkpoint_dir,
-                f"optim/rank{get_global_rank()}.pt",
-                optim.state_dict(),
-                upload_to=upload_to,
-                save_overwrite=self.cfg.save_overwrite,
-            )
-
-            # Save trainer state.
-            log.info("Saving trainer state...")
-            save_state_dict(
-                checkpoint_dir,
-                f"train/rank{get_global_rank()}.pt",
-                trainer_state,
-                upload_to=upload_to,
-                save_overwrite=self.cfg.save_overwrite,
-            )
-
-            # Save metadata.
-            self._save_metadata(checkpoint_dir, upload_to=upload_to)
-
-            # Save config. We do this last b/c the presence of a config in a remote checkpoint
-            # "directory" indicates that the folder is valid, as a opposed to a partially
-            # uploaded checkpoint directory that failed before completing.
-            self._save_config(checkpoint_dir, upload_to=upload_to)
-
-    def restore_checkpoint(
-        self,
-        load_path: PathOrStr,
-        fsdp_model: FSDP,
-        optim: Optimizer,
-        *,
-        local_cache: Optional[PathOrStr] = None,
-        load_optimizer_state: bool = True,
-    ) -> Dict[str, Any]:
-        # Load metadata and make sure checkpoint is compatible.
-        metadata = self._load_metadata(load_path, local_cache=local_cache)
-        assert metadata.world_size == get_world_size()
-
-        # Load local FSDP flat param data.
-        log.info("Loading local FSDP flat params data...")
-        model_state = load_state_dict(
-            load_path, f"model/rank{get_global_rank()}.pt", local_cache=local_cache, map_location="cpu"
-        )
-        self._load_flat_param_state(fsdp_model, model_state)
-        del model_state
-
-        # Load local optim state.
-        if load_optimizer_state:
-            log.info("Loading local optimizer state...")
-            optim_state = load_state_dict(
-                load_path, f"optim/rank{get_global_rank()}.pt", local_cache=local_cache, map_location="cpu"
-            )
-            # HACK/TODO (epwalsh): When we use adaptive clipping we track the 'grad_norm_exp_avg' for every param
-            # in every rank, and keep this in the optimizer state. But this causes issues when loading the
-            # state since torch sees the state is non-empty for some params which would normally be empty,
-            # and then assumes it should have all of the other state tensors for that param, which is doesn't.
-            # So for now we just remove 'grad_norm_exp_avg' everywhere from the state, which resets that metric.
-            # Not the end of the world but there's probably a better way around this without resetting
-            # the metric.
-            for param_id in list(optim_state["state"].keys()):
-                state = optim_state["state"][param_id]
-                if "grad_norm_exp_avg" in state:
-                    del state["grad_norm_exp_avg"]
-                if len(state) == 0:
-                    del optim_state["state"][param_id]
-            optim.load_state_dict(optim_state)
-            del optim_state
-
-        # Load local trainer state.
-        log.info("Loading local trainer state...")
-        trainer_state = load_state_dict(load_path, f"train/rank{get_global_rank()}.pt", local_cache=local_cache)
-        barrier()
-        return trainer_state
-
-    def _iter_flat_param_shards(
-        self, model_state: Dict[str, Any]
-    ) -> Generator[Tuple[str, _FlatParamShard], None, None]:
-        for module_data in model_state["modules"]:
-            module_prefix = module_data["name"].replace("_fsdp_wrapped_module.", "")
-            for handle in module_data["handles"]:
-                flat_data = handle["flat_param.data"]
-                if (num_padding := handle["flat_param._shard_numel_padded"]) > 0:
-                    # If there's padding in the flat param it should be on the right.
-                    assert (flat_data[-num_padding:] == 0).all()
-                # NOTE: this changes depending on the torch version, but we don't do a version
-                # check since we might be trying to unshard an old checkpoint that was stored
-                # with a different torch version than we're currently running with.
-                if "flat_param._shard_indices" in handle:
-                    # torch <=2.0.1
-                    param_start = handle["flat_param._shard_indices"][0]
-                    current_flat_index = 0
-                    for relative_fqn, full_shape, (offset_start, offset_end) in zip(
-                        handle["flat_param._fqns"][param_start:],
-                        handle["flat_param._shapes"][param_start:],
-                        handle["flat_param._shard_param_offsets"],
-                    ):
-                        root_fqn = relative_fqn if not module_prefix else f"{module_prefix}.{relative_fqn}"
-                        numel_shard = offset_end - offset_start + 1
-                        flat_param_shard = _FlatParamShard(
-                            full_shape=full_shape,
-                            shard_offsets=(offset_start, offset_end),
-                            shard_data=flat_data[current_flat_index : current_flat_index + numel_shard],
-                        )
-                        current_flat_index += numel_shard
-                        yield root_fqn, flat_param_shard
-                else:
-                    # torch >=2.1.0
-                    for relative_fqn, full_shape, shard_param_info in zip(
-                        handle["flat_param._fqns"],
-                        handle["flat_param._shapes"],
-                        handle["flat_param._shard_param_infos"],
-                    ):
-                        if not shard_param_info.in_shard:
-                            continue
-                        root_fqn = relative_fqn if not module_prefix else f"{module_prefix}.{relative_fqn}"
-                        flat_param_shard = _FlatParamShard(
-                            full_shape=full_shape,
-                            shard_offsets=(
-                                shard_param_info.intra_param_start_idx,
-                                shard_param_info.intra_param_end_idx,
-                            ),
-                            shard_data=flat_data[
-                                shard_param_info.offset_in_shard : shard_param_info.offset_in_shard
-                                + shard_param_info.numel_in_shard
-                            ],
-                        )
-                        yield root_fqn, flat_param_shard
-
-    def unshard_checkpoint(
-        self,
-        load_path: PathOrStr,
-        *,
-        local_cache: Optional[PathOrStr] = None,
-        load_optimizer_state: bool = True,
-        load_trainer_state: bool = True,
-        device: Optional[torch.device] = None,
-    ) -> Tuple[Dict[str, torch.Tensor], Optional[Dict[str, Any]], Optional[Dict[str, Any]]]:
-        device = device or torch.device("cpu")
-        metadata = self._load_metadata(load_path, local_cache=local_cache)
-
-        # Gather paths model state, potentially downloading them.
-        log.info("Gathering model state dicts...")
-        model_state_paths = self._gather_state_dict_paths(
-            load_path, "model", metadata.world_size, local_cache=local_cache
-        )
-
-        # Load model state dicts one-by-one, materializing and populating the full parameters as we go.
-        log.info("Materializing full parameters...")
-        full_model_state: Dict[str, torch.Tensor] = {}
-        # We keep a copy of the flat param metadata minus the actual tensors so we can reconstruct
-        # the full optimizer state below without having to reload the model state dicts.
-        flat_params_data: Dict[int, Dict[str, _FlatParamShard]] = defaultdict(dict)
-        for rank, path in enumerate(model_state_paths):
-            log.info(f"Loading shards from rank {rank}...")
-            model_state = torch.load(path, map_location="cpu")
-            for root_fqn, flat_param_shard in self._iter_flat_param_shards(model_state):
-                if root_fqn not in full_model_state:
-                    log.info(
-                        f"Materializing full parameter '{root_fqn}' with shape {flat_param_shard.full_shape}..."
-                    )
-                    assert flat_param_shard.shard_data is not None
-                    full_model_state[root_fqn] = torch.empty(
-                        flat_param_shard.full_shape, dtype=flat_param_shard.shard_data.dtype, device=device
-                    )
-                    # Fill with NaNs so we can validate that the whole parameter has been populated
-                    # afterwards.
-                    full_model_state[root_fqn].fill_(torch.nan)
-                # Copy over the local shard to the relevant part of the full parameter.
-                full_param = full_model_state[root_fqn]
-                log.info(f"Loading rank {rank} shard for '{root_fqn}'...")
-                flat_param_shard.copy_into(full_param)
-                flat_params_data[rank][root_fqn] = replace(flat_param_shard, shard_data=None)
-
-        log.info("Validating full parameters...")
-        for key, tensor in full_model_state.items():
-            if torch.isnan(tensor).any():
-                raise ValueError(f"Parameter '{key}' contains NaNs, this is likely a bug with the unsharder")
-
-        trainer_state: Optional[Dict[str, Any]] = None
-        if load_trainer_state:
-            trainer_state = load_state_dict(load_path, "train/rank0.pt", local_cache=local_cache)
-
-        if not load_optimizer_state:
-            return full_model_state, None, trainer_state
-
-        log.info("Gathering optim state dicts...")
-        optim_state_paths = self._gather_state_dict_paths(
-            load_path, "optim", metadata.world_size, local_cache=local_cache
-        )
-
-        log.info("Materializing full optim state...")
-        full_optim_state: Dict[str, Any] = {"state": defaultdict(dict)}
-        fqn_to_id: Dict[str, int] = {}
-        id_to_fqn: Dict[int, str] = {}
-        for rank, path in enumerate(optim_state_paths):
-            log.info(f"Loading sharded optim state from rank {rank}...")
-            optim_state = torch.load(path, map_location="cpu")
-
-            # Initialize param groups.
-            # We assume parameter groups are the same across all ranks.
-            # The only thing that differs across ranks is the state for each local sharded param.
-            if "param_groups" not in full_optim_state:
-                full_optim_state["param_groups"] = optim_state["param_groups"]
-            else:
-                assert full_optim_state["param_groups"] == optim_state["param_groups"]
-
-            # Generate mapping of parameter FQNs to optimizer param IDs and vice-versa.
-            if not fqn_to_id or not id_to_fqn:
-                for group in full_optim_state["param_groups"]:
-                    for fqn, id in zip(group["param_names"], group["params"]):
-                        fqn = fqn.replace("_fsdp_wrapped_module.", "")
-                        fqn_to_id[fqn] = id
-                        id_to_fqn[id] = fqn
-
-            # Iterate over local shard state and copy into the full state.
-            for id, shard_state in optim_state["state"].items():
-                fqn = id_to_fqn[id]
-                flat_param_shard = flat_params_data[rank].get(fqn)  # type: ignore[assignment]
-                full_state = full_optim_state["state"][id]
-                for key, shard_value in shard_state.items():
-                    assert isinstance(shard_value, torch.Tensor)
-                    if shard_value.shape == torch.Size([]):
-                        # Add singleton tensors directly to full state. These should be the same across
-                        # all ranks.
-                        assert key in ("step", "grad_norm_exp_avg")  # sanity check
-                        if key not in full_state:
-                            full_state[key] = shard_value.to(device)
-                        else:
-                            assert full_state[key] == shard_value
-                    else:
-                        # Otherwise we have a sharded param state.
-                        # If the corresponding full param state hasn't been materialized yet, do so now.
-                        assert flat_param_shard is not None, f"missing flat_params_data for {fqn} from rank {rank}"
-                        if key not in full_state:
-                            log.info(
-                                f"Materializing full state '{key}' for '{fqn}' with shape {flat_param_shard.full_shape}..."
-                            )
-                            full_state[key] = torch.empty(
-                                flat_param_shard.full_shape, dtype=shard_value.dtype, device=device
-                            )
-                        full_state_value = full_state[key]
-
-                        # Copy over the local shard state to the relevant part of the full parameter state.
-                        log.info(f"Loading rank {rank} shard state of '{key}' for '{fqn}'...")
-                        replace(flat_param_shard, shard_data=shard_value).copy_into(full_state_value)
-
-        # Lastly, clean up the parameter names in param groups.
-        for group in full_optim_state["param_groups"]:
-            group["param_names"] = [n.replace("_fsdp_wrapped_module.", "") for n in group["param_names"]]
-
-        return full_model_state, full_optim_state, trainer_state
-
-    def _get_state_dict_path(
-        self,
-        load_path: PathOrStr,
-        state_dict_type: str,
-        rank: int,
-        *,
-        local_cache: Optional[PathOrStr] = None,
-        progress=None,
-    ) -> Tuple[int, Path]:
-        fname = f"{state_dict_type}/rank{rank}.pt"
-        return rank, resource_path(str(load_path).rstrip("/"), fname, local_cache=local_cache, progress=progress)
-
-    def _gather_state_dict_paths(
-        self,
-        load_path: PathOrStr,
-        state_dict_type: str,
-        world_size: int,
-        *,
-        local_cache: Optional[PathOrStr] = None,
-    ) -> List[Path]:
-        progress = get_progress_bar()
-        with ThreadPoolExecutor(max_workers=self.thread_count) as executor:
-            futures = []
-            for rank in range(world_size):
-                future = executor.submit(
-                    self._get_state_dict_path,
-                    load_path,
-                    state_dict_type,
-                    rank,
-                    local_cache=local_cache,
-                    progress=progress,
-                )
-                futures.append(future)
-
-            results: Dict[int, Path] = {}
-            for future in as_completed(futures):
-                rank, path = future.result()
-                results[rank] = path
-
-        return [results[rank] for rank in range(world_size)]
-
-
-def build_sharded_checkpointer(
-    cfg: TrainConfig, *, name: Optional[ShardedCheckpointerType] = None
-) -> Checkpointer:
-    name = name or cfg.sharded_checkpointer
-    if name == ShardedCheckpointerType.torch_new:
-        return TorchNewStyleShardedCheckpointer(cfg)
-    elif name == ShardedCheckpointerType.torch_legacy:
-        return TorchLegacyShardedCheckpointer(cfg)
-    elif name == ShardedCheckpointerType.local:
-        return LocalShardedCheckpointer(cfg)
-    else:
-        raise NotImplementedError(name)

olmo_bitnet_1b/config.json

@@ -1,50 +0,0 @@
-{
-  "activation_type": "swiglu",
-  "alibi": false,
-  "alibi_bias_max": 8.0,
-  "architectures": [
-    "OLMoModelForCausalLM"
-  ],
-  "attention_dropout": 0.0,
-  "attention_layer_norm": false,
-  "attention_layer_norm_with_affine": false,
-  "bias_for_layer_norm": false,
-  "block_group_size": 1,
-  "block_type": "sequential",
-  "clip_qkv": null,
-  "d_model": 2048,
-  "embedding_dropout": 0.0,
-  "embedding_size": 50304,
-  "eos_token_id": 50279,
-  "flash_attention": true,
-  "include_bias": false,
-  "init_cutoff_factor": null,
-  "init_device": "cpu",
-  "init_fn": "mitchell",
-  "init_std": 0.02,
-  "layer_norm_type": "rms",
-  "layer_norm_with_affine": true,
-  "max_sequence_length": 2048,
-  "mlp_hidden_size": null,
-  "mlp_ratio": 8,
-  "model_type": "olmo",
-  "multi_query_attention": false,
-  "n_heads": 16,
-  "n_layers": 16,
-  "pad_token_id": 1,
-  "precision": "amp_bf16",
-  "residual_dropout": 0.0,
-  "rope": true,
-  "rope_full_precision": true,
-  "scale_logits": false,
-  "ternary": true,
-  "transformers_version": "4.38.2",
-  "use_cache": true,
-  "vocab_size": 50280,
-  "inference_mode":false,
-  "weight_tying": true,
-  "auto_map": {
-    "AutoConfig": "configuration_olmo.OLMoConfig",
-    "AutoModelForCausalLM": "modeling_olmo.OLMoForCausalLM"
-  }
-}

olmo_bitnet_1b/config.py

@@ -1,1106 +0,0 @@
-from __future__ import annotations
-
-from dataclasses import asdict, dataclass, field
-from glob import glob
-from pathlib import Path
-from typing import (
-    Any,
-    Dict,
-    Iterable,
-    List,
-    Optional,
-    Tuple,
-    Type,
-    TypeVar,
-    Union,
-    cast,
-)
-
-import torch
-from omegaconf import DictConfig, ListConfig
-from omegaconf import OmegaConf as om
-from omegaconf.errors import OmegaConfBaseException
-from torch.distributed.fsdp import MixedPrecision, ShardingStrategy
-
-from .aliases import PathOrStr
-from .beam_search import Sampler
-from .exceptions import OLMoConfigurationError
-from .util import StrEnum
-
-__all__ = [
-    "ActivationType",
-    "ActivationCheckpointingStrategy",
-    "BlockType",
-    "LayerNormType",
-    "InitFnType",
-    "ModelConfig",
-    "OptimizerType",
-    "OptimizerConfig",
-    "SchedulerType",
-    "SchedulerConfig",
-    "DataConfig",
-    "EvaluatorConfig",
-    "TokenizerConfig",
-    "TrainConfig",
-    "PaddingDirection",
-    "TruncationDirection",
-    "SpeedMonitorConfig",
-    "WandbConfig",
-    "CompilerConfig",
-    "WandbConfig",
-    "FSDPPrecision",
-    "FSDPWrapStrategy",
-    "FSDPConfig",
-    "CheckpointType",
-]
-
-C = TypeVar("C", bound="BaseConfig")
-D = TypeVar("D", bound="DictConfig|ListConfig")
-
-
-class BaseConfig:
-    @classmethod
-    def _register_resolvers(cls, validate_paths: bool = True):
-        # Expands path globs into a list.
-        def path_glob(*paths) -> List[str]:
-            out = []
-            for path in paths:
-                matches = sorted(glob(path))
-                if not matches and validate_paths:
-                    raise FileNotFoundError(f"{path} does not match any files or dirs")
-                out.extend(matches)
-            return out
-
-        # Chooses the first path in the arguments that exists.
-        def path_choose(*paths) -> str:
-            from .util import is_url
-
-            for path in paths:
-                if is_url(path) or Path(path).exists():
-                    return path
-            if validate_paths:
-                raise FileNotFoundError(", ".join(paths))
-            else:
-                return ""
-
-        # Finds the latest checkpoint in a folder.
-        def path_last_checkpoint(path) -> str:
-            from .util import find_latest_checkpoint
-
-            latest_checkpoint = find_latest_checkpoint(path)
-            if latest_checkpoint is None:
-                if validate_paths:
-                    raise FileNotFoundError(f"Could not find a latest checkpoint at {path}")
-                else:
-                    return ""
-            else:
-                return str(latest_checkpoint)
-
-        om.register_new_resolver("path.glob", path_glob, replace=True)
-        om.register_new_resolver("path.choose", path_choose, replace=True)
-        om.register_new_resolver("path.last_checkpoint", path_last_checkpoint, replace=True)
-
-    @classmethod
-    def update_legacy_settings(cls, config: D) -> D:
-        """
-        Update the legacy config settings whose schemas have undergone backwards-incompatible changes.
-        """
-        return config
-
-    @classmethod
-    def new(cls: Type[C], **kwargs) -> C:
-        cls._register_resolvers()
-        conf = om.structured(cls)
-        try:
-            if kwargs:
-                conf = om.merge(conf, kwargs)
-            return cast(C, om.to_object(conf))
-        except OmegaConfBaseException as e:
-            raise OLMoConfigurationError(str(e))
-
-    @classmethod
-    def load(
-        cls: Type[C],
-        path: PathOrStr,
-        overrides: Optional[List[str]] = None,
-        key: Optional[str] = None,
-        validate_paths: bool = True,
-    ) -> C:
-        """Load from a YAML file."""
-        cls._register_resolvers(validate_paths=validate_paths)
-        schema = om.structured(cls)
-        try:
-            raw = om.load(str(path))
-            if key is not None:
-                raw = raw[key]  # type: ignore
-            raw = cls.update_legacy_settings(raw)
-            conf = om.merge(schema, raw)
-            if overrides:
-                conf = om.merge(conf, om.from_dotlist(overrides))
-            return cast(C, om.to_object(conf))
-        except OmegaConfBaseException as e:
-            raise OLMoConfigurationError(str(e))
-
-    def save(self, path: PathOrStr) -> None:
-        """Save to a YAML file."""
-        om.save(config=self, f=str(path))
-
-    def asdict(self, exclude: Optional[Iterable[str]] = None) -> Dict[str, Any]:
-        out = asdict(self)  # type: ignore
-        if exclude is not None:
-            for name in exclude:
-                if name in out:
-                    del out[name]
-        return out
-
-
-class LayerNormType(StrEnum):
-    default = "default"
-    """
-    The default LayerNorm implementation, equivalent to PyTorch's built-in version.
-    """
-
-    low_precision = "low_precision"
-    """
-    A low-precision version of the default LayerNorm.
-    """
-
-    rms = "rms"
-    """
-    An RMSNorm implementation. When using ``torch.compile`` this is
-    probably the fastest implementation.
-    """
-
-
-class ActivationType(StrEnum):
-    gelu = "gelu"
-    relu = "relu"
-    swiglu = "swiglu"
-
-
-class BlockType(StrEnum):
-    sequential = "sequential"
-
-    llama = "llama"
-    """
-    A block similar to the sequential block with slightly different
-    implementations of operations like attention to imitate the behavior of Llama.
-    """
-
-
-class InitFnType(StrEnum):
-    mitchell = "mitchell"
-    """
-    The strategy suggested to us by Mitchell Wortsman from UW.
-    This uses a truncated normal distribution with an adaptive standard deviation that depends
-    on the size of the weights as well as the depth of the layer.
-    """
-
-    normal = "normal"
-    """
-    All weights are initialized from the same normal distribution.
-    """
-
-    kaiming_normal = "kaiming_normal"
-    """
-    All weights are initialized with the Kaiming method from a normal distribution.
-    Note this currently won't work with FSDP.
-    """
-
-    fan_in = "fan_in"
-    """
-    "Fan-in variance scaling", i.e. normal with a standard deviation of ``1/sqrt(d_in)`` where ``d_in``
-    is the input dimensionality of the kernel.
-    """
-
-    full_megatron = "full_megatron"
-    """
-    This is what metaseq calls "full megatron init". It is the init used for Llama 2.
-    """
-
-
-@dataclass
-class ModelConfig(BaseConfig):
-    """
-    OLMo (model) configuration.
-    """
-
-    # Note that the defaults for these attributes are equivalent to the base GPT2 model.
-
-    d_model: int = 768
-    """
-    The hidden size of the model.
-    """
-
-    n_heads: int = 12
-    """
-    The number of self-attention heads.
-    """
-
-    n_kv_heads: Optional[int] = None
-    """
-    The number of heads to use for keys and values. Defaults to `n_heads`.
-    Set this to ``None`` or ``n_heads`` for normal multi-head attention.
-    Set this to 1 for multi-query attention.
-    Set it to some in-between value for Llama2-style grouped query attention.
-    """
-
-    clip_qkv: Optional[float] = None
-    """
-    Clip QKV to this value when set.
-    """
-
-    n_layers: int = 12
-    """
-    The number of layers/blocks.
-    """
-
-    mlp_ratio: int = 4
-    """
-    The ratio of the inner MLP dimensionality to ``d_model``.
-    This is only used when ``mlp_hidden_size`` is not set.
-    """
-
-    mlp_hidden_size: Optional[int] = None
-    """
-    Set the exact hidden size for the MLP. Otherwise the inner MLP hidden size will be set to `mlp_ratio * d_model`.
-    """
-
-    activation_type: ActivationType = ActivationType.swiglu
-    """
-    The activation function to use within the MLP layers.
-    """
-
-    block_type: BlockType = BlockType.sequential
-    """
-    The transformer block implementation.
-    """
-
-    block_group_size: int = 1
-    """
-    The number of blocks to group together into a single parent block.
-    This has no affect on the number of parameters in the model and is only used to wrap groups
-    of blocks together with a single FSDP wrapper during training.
-    """
-
-    alibi: bool = False
-    """
-    If ``True``, use ALiBi embeddings. Mutually exclusive with ``rope``.
-    """
-
-    alibi_bias_max: float = 8.0
-    """
-    Maximum absolute value of ALiBi bias.
-    """
-
-    rope: bool = False
-    """
-    Use rotary positional embeddings (RoPE). Mutually exclusive with ``alibi``.
-    """
-
-    rope_full_precision: bool = True
-    """
-    If ``True``, apply RoPE embeddings at full precision regardless of the input type. Otherwise,
-    apply RoPE at the precision of the input.
-    """
-
-    flash_attention: bool = False
-    """
-    If ``True``, use ``FlashAttention``.
-    """
-
-    attention_dropout: float = 0.1
-    """
-    The dropout probability within the attention modules.
-    """
-
-    multi_query_attention: Optional[bool] = None
-    """
-    Deprecated. Use n_kv_heads instead.
-    """
-
-    attention_layer_norm: bool = False
-    """
-    Apply layer norm to the keys and queries within the attention mechanism.
-    This can help stabilize training.
-    """
-
-    residual_dropout: float = 0.1
-    """
-    The dropout probability for the MLP and attention output within each block.
-    """
-
-    embedding_dropout: float = 0.1
-    """
-    The dropout probability for embeddings.
-    """
-
-    layer_norm_type: LayerNormType = LayerNormType.default
-    """
-    The layernorm implementation to use.
-    """
-
-    layer_norm_with_affine: bool = True
-    """
-    Whether to include bias and weight parameters for the layer norms.
-    This only affects layer norms that are immediately followed by a linear layer in the forward pass,
-    so everything except QK-norms. To turn off affines for QK norms as well, set :attr:`attention_layer_norm_with_affine`
-    to ``False``.
-    """
-
-    attention_layer_norm_with_affine: bool = True
-    """
-    Toggle affine transform for the QK norms.
-    """
-
-    max_sequence_length: int = 1024
-    """
-    The maximum input sequence length supported by the model.
-    """
-
-    include_bias: bool = True
-    """
-    Whether or not to include bias parameters in linear layers.
-    In PaLM, they got rid of all bias terms because they found that large
-    models tend to have near 0 bias terms anyway.
-    """
-
-    bias_for_layer_norm: Optional[bool] = None
-    """
-    Whether or not to include bias parameters in layer norm.
-    This is separate from the include_bias parameter, because of a ROCm crash when biases are disabled in
-    layer norm.
-    When this is None (the default), it inherits the setting from include_bias.
-    """
-
-    scale_logits: bool = False
-    """
-    If ``True``, scale the output logits by ``1 / sqrt(d_model)``.
-    """
-
-    vocab_size: int = 50257
-    """
-    Vocabulary size of the model.
-    """
-
-    embedding_size: Optional[int] = 50304
-    """
-    The number of embeddings, i.e. the number of tokens. If set to ``None`` it will default
-    to ``vocab_size``. If ``vocab_size`` is not a multiple of 128, setting this to the
-    next multiple of 128 that's greater than ``vocab_size`` can improve throughput
-    substantially.
-    """
-
-    weight_tying: bool = True
-    """
-    Whether to tie output linear weights to the input embedding.
-    """
-
-    eos_token_id: int = 50256
-    """
-    The ID of the end-of-sentence special token.
-    """
-
-    pad_token_id: int = 50256
-    """
-    The ID of the token to use for padding. Defaults to the ID of the EOS token.
-    """
-
-    init_device: Optional[str] = None
-    """
-    The torch device to use when initializing the model parameters, e.g. "cpu", "cuda:0", "meta".
-    """
-
-    init_fn: InitFnType = InitFnType.normal
-    """
-    The weight initialization strategy.
-    """
-
-    init_std: float = 0.02
-    """
-    The standard deviation to use when initializing weights with a "fixed distribution" ``init_fn``, such
-    as "normal".
-    """
-
-    init_cutoff_factor: Optional[float] = None
-    """
-    A positive factor used to scale the cutoff values when initializing weights with a "fixed distribution" ``init_fn``, such
-    as "normal". Setting this to None means values are not cutoff.
-    """
-
-    precision: Optional[str] = None
-    """
-    Precision used to train/evaluate with. You shouldn't set this directly.
-    See :data:`TrainConfig.precision` instead.
-    """
-
-    ternary: bool = False
-    """
-    Use ternary BitLinear layer from "The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits" (https://arxiv.org/pdf/2402.17764.pdf)
-    """
-
-    @property
-    def effective_n_kv_heads(self) -> int:
-        if self.n_kv_heads is None:
-            if self.multi_query_attention is True:
-                return 1
-            else:
-                return self.n_heads
-        else:
-            if self.multi_query_attention is None:
-                return self.n_kv_heads
-            if self.multi_query_attention:
-                n_kv_heads_should_be = 1
-            else:
-                n_kv_heads_should_be = self.n_heads
-            if self.n_kv_heads == n_kv_heads_should_be:
-                return n_kv_heads_should_be
-            else:
-                raise OLMoConfigurationError(
-                    "You can't set `multi_query_attention` and `n_kv_heads` at the same time."
-                )
-
-
-class OptimizerType(StrEnum):
-    lionw = "lionw"
-    adamw = "adamw"
-
-
-@dataclass
-class OptimizerConfig(BaseConfig):
-    name: OptimizerType = OptimizerType.lionw
-    learning_rate: float = 1.0e-4
-    weight_decay: float = 0.01
-    betas: Tuple[float, float] = (0.9, 0.95)
-
-    no_decay_norm_and_bias: Optional[bool] = None
-    """
-    Deprecated. Use ``decay_norm_and_bias`` and ``decay_embeddings`` instead.
-    """
-
-    decay_norm_and_bias: bool = False
-    decay_embeddings: bool = False
-    metrics_log_interval: Optional[int] = None
-    """
-    The interval with which to collect and log detailed parameter-specific metrics.
-    This only applies when logging to W&B, since these metrics won't be logged to the console.
-    If not set, defaults to the wandb `log_interval`.
-    """
-
-    def __post_init__(self):
-        self.betas = tuple(self.betas)  # type: ignore[assignment]
-
-    @classmethod
-    def update_legacy_settings(cls, config: D) -> D:
-        new_config = config.copy()
-        if om.is_dict(new_config):
-            assert isinstance(new_config, DictConfig)
-
-            if hasattr(new_config, "name") and new_config.name == "decoupled_lionw":
-                new_config.name = "lionw"
-                if hasattr(new_config, "eps"):
-                    del new_config.eps
-
-        return new_config
-
-
-class SchedulerType(StrEnum):
-    cosine_with_warmup = "cosine_with_warmup"
-    linear_with_warmup = "linear_with_warmup"
-    inverse_sqrt_with_warmup = "inverse_sqrt_with_warmup"
-    max_scheduler = "max_scheduler"
-    constant = "constant"
-
-
-class SchedulerUnits(StrEnum):
-    steps = "steps"
-    tokens = "tokens"
-
-
-@dataclass
-class SchedulerConfig(BaseConfig):
-    name: SchedulerType = SchedulerType.cosine_with_warmup
-    units: SchedulerUnits = SchedulerUnits.steps
-    t_warmup: Union[int, float] = 100
-    t_max: Optional[Union[int, float]] = None
-    alpha_f: float = 0.1
-
-    grad_clip_warmup_steps: Optional[Union[int, float]] = None
-    """
-    The warmup period for which the max grad norm (or norm ratio) will be set to its
-    warmup value of `max_grad_norm * grad_clip_warmup_factor`.
-    """
-
-    grad_clip_warmup_factor: Optional[float] = None
-    """
-    The ratio of the max allowed gradient norm (or norm ratio) for clipping during the warmup period
-    vs after the warmup period.
-    """
-
-
-class PaddingDirection(StrEnum):
-    right = "right"
-    left = "left"
-
-
-@dataclass
-class DataConfig(BaseConfig):
-    paths: Optional[List[str]] = None
-    datasets: Optional[Dict[str, List[str]]] = None
-    label_mask_paths: Optional[List[str]] = None
-    pad_direction: PaddingDirection = PaddingDirection.right
-    generate_attention_mask: bool = False
-    num_workers: int = 0
-    drop_last: bool = False
-    pin_memory: bool = False
-    prefetch_factor: Optional[int] = None
-    persistent_workers: bool = False
-    timeout: int = 0
-    seed: Optional[int] = None
-
-
-class EvaluatorType(StrEnum):
-    downstream = "downstream"
-    lm = "lm"
-
-
-@dataclass
-class EvaluatorConfig(BaseConfig):
-    label: str
-    type: EvaluatorType = EvaluatorType.lm
-    data: DataConfig = field(default_factory=DataConfig)
-    device_eval_batch_size: Optional[int] = None
-    subset_num_batches: Optional[int] = None
-
-
-class TruncationDirection(StrEnum):
-    right = "right"
-    left = "left"
-
-
-@dataclass
-class TokenizerConfig(BaseConfig):
-    identifier: str = "gpt2"
-    truncate_direction: TruncationDirection = TruncationDirection.right
-
-
-@dataclass
-class WandbConfig(BaseConfig):
-    project: Optional[str] = None
-    entity: Optional[str] = "ai2-llm"
-    group: Optional[str] = None
-    name: Optional[str] = None
-    tags: Optional[List[str]] = field(default_factory=lambda: ["watching"])
-    log_artifacts: bool = False
-    rank_zero_only: bool = True
-    log_interval: int = 1
-
-
-@dataclass
-class SpeedMonitorConfig(BaseConfig):
-    window_size: int = 100
-    gpu_flops_available: Optional[Union[float, int]] = None
-
-
-@dataclass
-class CompilerConfig(BaseConfig):
-    mode: Optional[str] = None
-    """
-    The mode to compile the model in. At the moment this can be "default",
-    "reduce-overhead" (useful for smaller models/batches), or "max-autotune"
-    (the fastest for larger models, but takes a long time to compile).
-    """
-
-    fullgraph: bool = False
-    """
-    Whether it is OK to break model into several subgraphs when compiling.
-    Note that this is not compatible with FSDP.
-    """
-
-    backend: str = "inductor"
-    """
-    The backend to use.
-    """
-
-
-class FSDPWrapStrategy(StrEnum):
-    by_block = "by_block"
-    """
-    Wrap each OLMo block with its own FSDP instance.
-    """
-
-    by_block_and_size = "by_block_and_size"
-    """
-    Like 'by_block' but `wte` and `ff_out` will be wrapped separately as well.
-    """
-
-    by_block_group = "by_block_group"
-    """
-    Wrap each block group together into its own FSDP instance.
-    This requires :attr:`~ModelConfig.block_group_size` to be bigger than 1.
-    """
-
-    by_block_group_and_size = "by_block_group_and_size"
-    """
-    Like 'by_block_group' but `wte` and `ff_out` will be wrapped separately as well.
-    """
-
-    size_based = "size_based"
-    """
-    Used PyTorch's default size-based auto wrap policy.
-    """
-
-    one_in_two = "one_in_two"
-    one_in_three = "one_in_three"
-    one_in_four = "one_in_four"
-    one_in_five = "one_in_five"
-
-
-class FSDPPrecision(StrEnum):
-    pure = "pure"
-    """
-    Equivalent to :class:`torch.distributed.fsdp.MixedPrecision` with ``param_dtype``, ``reduce_dtype``,
-    and ``buffer_dtype`` all set to the autocast precision data type.
-    """
-
-    mixed = "mixed"
-    """
-    Equivalent to :class:`torch.distributed.fsdp.MixedPrecision` with ``param_dtype``, and ``buffer_dtype``
-    set to the autocast precision data type, while ``reduce_dtype`` is set to fp32.
-    """
-
-
-@dataclass
-class FSDPConfig(BaseConfig):
-    use_orig_params: bool = True
-    """
-    This must be ``True`` if using ``compile`` or you want to track the parameter norm during training.
-    """
-
-    sharding_strategy: ShardingStrategy = ShardingStrategy.FULL_SHARD
-
-    wrapping_strategy: Optional[FSDPWrapStrategy] = None
-    """
-    The wrapping strategy to use. If ``None``, the default, the model is wrapped with a single top-level
-    FSDP instance.
-    """
-
-    precision: FSDPPrecision = FSDPPrecision.pure
-
-
-class CheckpointType(StrEnum):
-    sharded = "sharded"
-    unsharded = "unsharded"
-    sharded_ephemeral = "sharded_ephemeral"
-
-
-class ShardedCheckpointerType(StrEnum):
-    torch_new = "torch_new"
-    torch_legacy = "torch_legacy"
-    local = "local"
-
-
-class ActivationCheckpointingStrategy(StrEnum):
-    whole_layer = "whole_layer"
-    """
-    Checkpoint every transformer layer.
-    """
-
-    one_in_two = "one_in_two"
-    """
-    Checkpoint one in two transformer layers.
-    """
-
-    one_in_three = "one_in_three"
-    """
-    Checkpoint one in three transformer layers.
-    """
-
-    one_in_four = "one_in_four"
-    """
-    Checkpoint one in four transformer layers.
-    """
-
-    two_in_three = "two_in_three"
-    """
-    Checkpoint two out of every three transformer layers.
-    """
-
-    three_in_four = "three_in_four"
-    """
-    Checkpoint three out of four of every transformer layers.
-    """
-
-    fine_grained = "fine_grained"
-    """
-    Focus checkpointing on where it is cheap to recompute and saves most memory.
-    """
-
-
-@dataclass
-class TrainConfig(BaseConfig):
-    """
-    OLMo training configuration.
-    """
-
-    run_name: Optional[str] = None
-    """
-    The name of the run.
-    """
-
-    seed: int = 6198
-    """
-    Used to seed all initial RNG states.
-    """
-
-    epoch: Optional[int] = None
-    """
-    Increment this when starting a new epoch.
-    """
-
-    dry_run: bool = False
-    """
-    If ``True``, don't actually train.
-    """
-
-    model: ModelConfig = field(default_factory=ModelConfig)
-    """
-    OLMo Model configuration.
-    """
-
-    optimizer: OptimizerConfig = field(default_factory=OptimizerConfig)
-    """
-    Optimizer configuration.
-    """
-
-    scheduler: SchedulerConfig = field(default_factory=SchedulerConfig)
-    """
-    Learning rate scheduler configuration.
-    """
-
-    data: DataConfig = field(default_factory=DataConfig)
-    """
-    Training data configuration.
-    """
-
-    restore_dataloader: bool = True
-    """
-    When restarting, restore the data loader to where it left off.
-    If you restarting in order to train on a different dataset, set this to ``False``.
-    """
-
-    fast_forward_batches: Optional[int] = None
-    """
-    When restarting, use this to fast-forward the dataloader beyond the last checkpoint.
-    This can be useful when restarting due to a loss spike in order to skip the data that
-    corresponded to the spike.
-    """
-
-    evaluators: List[EvaluatorConfig] = field(default_factory=list)
-    """
-    Evaluation configurations.
-    """
-
-    eval_interval: int = 1000
-    """
-    How often (in terms of batches) to run evaluations.
-    """
-
-    tokenizer: TokenizerConfig = field(default_factory=TokenizerConfig)
-    """
-    Tokenizer configuration.
-    """
-
-    save_folder: str = "./"
-    """
-    The directory to save checkpoints to.
-    """
-
-    remote_save_folder: Optional[str] = None
-    """
-    A folder in a cloud bucket to upload saved checkpoints to.
-    """
-
-    canceled_check_interval: int = 50
-    """
-    How often (in batches) to check if the run has been canceled or reached its time limit.
-    """
-
-    save_interval: int = 1000
-    """
-    How often (in terms of steps) to save sharded training state checkpoints.
-    """
-
-    save_interval_unsharded: Optional[int] = None
-    """
-    How often (if at all) to save unsharded training state checkpoint.
-    For large models it can be costly to save these, so it usually makes sense to save
-    these less often than regular (sharded) training checkpoints.
-    """
-
-    save_interval_ephemeral: Optional[int] = None
-    """
-    How often (if at all) to save ephemeral sharded checkpoints. These checkpoints are the same
-    as those saved every `save_interval` except that at most only the most recent one of these is kept.
-    This is useful when you want to checkpoint often for restarts in case of failures, but don't
-    want to keep the majority of these checkpoints.
-
-    For example, suppose you want to keep your checkpoints at every 1000 steps, but you also want to save
-    a temporary checkpoint every 100 steps in case your job fails. In that case you would
-    set `save_interval=1000` and `save_interval_ephemeral=100`.
-    """
-
-    save_num_checkpoints_to_keep: int = -1
-    """
-    How many sharded checkpoints to keep.
-    """
-
-    save_num_unsharded_checkpoints_to_keep: int = -1
-    """
-    How many unsharded checkpoints to keep.
-    """
-
-    save_overwrite: bool = False
-    """
-    If ``True``, overwrite any conflicting checkpoint files.
-    """
-
-    force_save_unsharded: bool = False
-    """
-    Save an unsharded checkpoint before training (even during a dry run).
-    Use this option with `--load-path={PATH}` and `--dry_run` to convert a sharded
-    checkpoint into an unsharded checkpoint.
-    """
-
-    no_pre_train_checkpoint: bool = False
-    """
-    Skip saving pre-train checkpoint.
-    """
-
-    load_path: Optional[str] = None
-    """
-    The path to a training checkpoint to restore/resume from.
-
-    Note that you can make use of the "path.last_checkpoint" Omegaconfig YAML resolver here, which takes
-    a local or remote directory and resolves to the latest checkpoint (sharded or unsharded) in that directory.
-    For example,
-
-    ```bash
-    --load_path='${path.last_checkpoint:s3://ai2-llm/checkpoints/7b/v1_5-mix-run-001}'
-    ```
-    """
-
-    load_path_sharded_checkpointer: Optional[ShardedCheckpointerType] = None
-    """
-    The sharded checkpointer type to use to load the initial checkpoint from ``load_path``.
-    """
-
-    reset_optimizer_state: bool = False
-    """
-    When this is set, we restore the model from a checkpoint (if given), but we leave the optimizer uninitialized.
-    We also set a new learning rate schedule that does a new warmup, such that it intercepts the original learning
-    curve (according to the current learning rate schedule settings), and continues from there.
-    """
-
-    reset_trainer_state: bool = False
-    """
-    When this is set we don't restore the trainer state from a checkpoint.
-    """
-
-    sharded_checkpointer: ShardedCheckpointerType = ShardedCheckpointerType.torch_legacy
-    """
-    The name of the sharded checkpointer to use to save (sharded) checkpoints throughout training.
-    """
-
-    new_style_checkpoints: Optional[bool] = None
-    """
-    Deprecated. Use ``sharded_checkpointer`` instead.
-    """
-
-    max_duration: Union[int, str] = 10000
-    """
-    How long to train for.
-
-    If specified without a unit (the default), the units are assumed to be steps.
-    You can also specify this in terms of tokens, for example: `max_duration="2e12T"` means train until
-    2 trillion tokens.
-    """
-
-    global_train_batch_size: int = 512
-    """
-    The effective global batch size.
-    """
-
-    device_train_batch_size: Optional[int] = None  # calculated automatically
-    """
-    Don't set this manually. This will be set to ``global_train_batch_size // world_size``.
-    """
-
-    device_train_microbatch_size: int = 16
-    """
-    The number of instances passed to the model in a single forward-backward pass. You should set
-    this as large as you can based on available GPU memory.
-    """
-
-    device_eval_batch_size: int = 16
-    """
-    The number of evaluation instances passed to the model in a single forward pass on each device.
-    """
-
-    eval_subset_num_batches: int = -1
-    """
-    The number of batches to use for downstream evaluation from each dataset.
-    """
-
-    eval_on_load: bool = False
-    """
-    When resuming from a checkpoint, run the evaluation loop right away.
-    """
-
-    device_train_grad_accum: Optional[int] = None  # calculated automatically
-    """
-    Don't set this manually. This will be set to ``device_train_batch_size // device_train_microbatch_size``.
-    """
-
-    max_grad_norm: Optional[float] = None
-    """
-    Clip gradient norms to this value if set.
-    """
-
-    max_grad_norm_ratio: Optional[float] = None
-    """
-    If set, gradient norms will be clipped to `max_grad_norm_ratio * exp_avg(norm(grad))`.
-    This takes priority over `max_grad_norm` when set.
-    """
-
-    precision: Optional[str] = None
-    """
-    Precision to train with (e.g. "amp_bf16", "amp_fp16", or "fp32").
-    """
-
-    wandb: Optional[WandbConfig] = None
-    """
-    Weights & Biases configuration.
-    """
-
-    speed_monitor: SpeedMonitorConfig = field(default_factory=SpeedMonitorConfig)
-    """
-    Speed monitor configuration.
-    """
-
-    console_log_interval: int = 1
-    """
-    How often to log to the console.
-    """
-
-    compile: Optional[CompilerConfig] = None
-    """
-    Settings for compiling the model with ``torch.compile()``.
-    """
-
-    fsdp: FSDPConfig = field(default_factory=FSDPConfig)
-    """
-    Fully sharded data parallel settings.
-    """
-
-    softmax_auxiliary_loss: bool = False
-    """
-    If ``True``, we add the auxiliary loss function from PaLM that encourages the softmax
-    normalizing term to be close to 0.
-    """
-
-    time_limit: Optional[float] = 60 * 60 * 47.5
-    """
-    The maximum amount of time to train for before saving a checkpoint and ending early.
-    On LUMI we have 48 hours max per job, so we default to just under 48 hours to give us time
-    to write out a final checkpoint.
-    """
-
-    extra_steps_after_cancel: int = 10
-    """
-    Under certain conditions when a run is canceled we train for a few extra steps after saving
-    the final checkpoint so that when the run is restarted from the latest checkpoint we have some
-    overlap in metrics.
-    """
-
-    early_stopping_factor: Optional[float] = None
-
-    save_data_indices: bool = True
-    """
-    Save training data indices from each batch for each worker.
-    """
-
-    python_profiling: bool = False
-    """
-    Whether to run the Python profiler on batches 6, 7, and 8.
-    """
-
-    torch_profiling: bool = False
-    """
-    Whether to run the PyTorch profiler on batches 6, 7, and 8.
-    """
-
-    stop_at: Optional[int] = None
-    """
-    Stop at a specific step.
-    """
-
-    stop_after: Optional[int] = None
-    """
-    Stop after a specific number of steps.
-    """
-
-    activation_checkpointing: Optional[ActivationCheckpointingStrategy] = None
-    """
-    The activation checkpointing strategy to use.
-    """
-
-    fused_loss: Optional[bool] = None
-    """
-    Whether to use the fused CE loss function from `flash-attn`.
-    """
-
-    @property
-    def autocast_precision(self) -> torch.dtype:
-        if self.precision == "amp_bf16":
-            return torch.bfloat16
-        elif self.precision == "amp_fp16":
-            return torch.float16
-        elif self.precision == "fp32":
-            return torch.float32
-        else:
-            raise ValueError(f"Unexpected precision type '{self.precision}'")
-
-    @property
-    def fsdp_precision(self) -> MixedPrecision:
-        if self.fsdp.precision == FSDPPrecision.pure:
-            return MixedPrecision(
-                param_dtype=self.autocast_precision,
-                reduce_dtype=self.autocast_precision,
-                buffer_dtype=self.autocast_precision,
-            )
-        elif self.fsdp.precision == FSDPPrecision.mixed:
-            return MixedPrecision(
-                param_dtype=self.autocast_precision,
-                reduce_dtype=torch.float32,
-                buffer_dtype=self.autocast_precision,
-            )
-        else:
-            raise NotImplementedError(f"{self.fsdp.precision}")
-
-    @classmethod
-    def update_legacy_settings(cls, config: D) -> D:
-        new_config = config.copy()
-        if om.is_dict(new_config):
-            assert isinstance(new_config, DictConfig)
-
-            if hasattr(new_config, "activation_checkpointing"):
-                if new_config.activation_checkpointing is False:
-                    new_config.activation_checkpointing = None
-                if new_config.activation_checkpointing is True:
-                    new_config.activation_checkpointing = ActivationCheckpointingStrategy.whole_layer
-
-            if hasattr(new_config, "optimizer"):
-                new_config.optimizer = OptimizerConfig.update_legacy_settings(new_config.optimizer)
-
-        return new_config

olmo_bitnet_1b/configuration_olmo.py

@@ -1,52 +0,0 @@
-"""
-OLMo configuration
-"""
-
-from transformers import AutoConfig, PretrainedConfig
-from transformers.utils import logging
-
-from .config import ModelConfig
-from .aliases import PathOrStr
-from .beam_search import Sampler
-from .exceptions import OLMoError
-from .initialization import ModuleType
-from .optim import Optimizer
-from .util import StrEnum
-from .safetensors_util import STKey
-from .torch_util import seed_all
-
-logger = logging.get_logger(__name__)
-
-
-class OLMoConfig(PretrainedConfig):
-    model_type = "olmo"
-    keys_to_ignore_at_inference = ["past_key_values"]  # TODO: confirm
-
-    def __init__(self, use_cache: bool = False, **kwargs):
-        model_config = ModelConfig()
-        all_kwargs = model_config.asdict()
-        all_kwargs.update(kwargs)
-        all_kwargs.update({"use_cache": use_cache})
-        all_kwargs.update(
-            {
-                "architectures": all_kwargs.get("architectures", ["OLMoModelForCausalLM"])
-                or ["OLMoModelForCausalLM"]
-            }
-        )
-        super().__init__(**all_kwargs)
-
-    @property
-    def num_attention_heads(self):
-        return self.n_heads
-
-    @property
-    def num_hidden_layers(self):
-        return self.n_layers
-
-    @property
-    def hidden_size(self):
-        return self.d_model
-
-
-# Register the config class so that it is available for transformer pipelines, auto-loading etc.
-# AutoConfig.register("olmo", OLMoConfig)

olmo_bitnet_1b/exceptions.py

@@ -1,50 +0,0 @@
-__all__ = [
-    "OLMoError",
-    "OLMoConfigurationError",
-    "OLMoCliError",
-    "OLMoEnvironmentError",
-    "OLMoNetworkError",
-    "OLMoCheckpointError",
-]
-
-
-class OLMoError(Exception):
-    """
-    Base class for all custom OLMo exceptions.
-    """
-
-
-class OLMoConfigurationError(OLMoError):
-    """
-    An error with a configuration file.
-    """
-
-
-class OLMoCliError(OLMoError):
-    """
-    An error from incorrect CLI usage.
-    """
-
-
-class OLMoEnvironmentError(OLMoError):
-    """
-    An error from incorrect environment variables.
-    """
-
-
-class OLMoNetworkError(OLMoError):
-    """
-    An error with a network request.
-    """
-
-
-class OLMoCheckpointError(OLMoError):
-    """
-    An error occurred reading or writing from a checkpoint.
-    """
-
-
-class OLMoThreadError(Exception):
-    """
-    Raised when a thread fails.
-    """

olmo_bitnet_1b/initialization.py

@@ -1,95 +0,0 @@
-import math
-from typing import Optional, Union
-
-import torch
-import torch.nn as nn
-
-from .config import InitFnType, ModelConfig
-from .util import StrEnum
-
-__all__ = ["init_weights", "ModuleType"]
-
-
-class ModuleType(StrEnum):
-    in_module = "in"
-    out_module = "out"
-    emb = "emb"
-    final_out = "final_out"
-
-
-def init_weights(
-    config: ModelConfig,
-    module: Union[nn.Linear, nn.Embedding],
-    d: Optional[int] = None,
-    layer_id: Optional[int] = None,
-    std_factor: float = 1.0,
-    type_of_module: Optional[ModuleType] = None,
-) -> None:
-    """
-    Initialize weights of a linear or embedding module.
-
-    :param config: The model config.
-    :param module: The linear or embedding submodule to initialize.
-    :param d: The effective input dimensionality of the weights. This could be smaller than the actual dimensions
-        for fused layers.
-    :param layer_id: When set, the standard deviation for the "mitchell" method will be adjusted by
-        ``1 / sqrt(2 * (layer_id + 1))``.
-    """
-    d = d if d is not None else config.d_model
-    if config.init_fn == InitFnType.normal:
-        std = config.init_std * std_factor
-        if config.init_cutoff_factor is not None:
-            cutoff_value = config.init_cutoff_factor * std
-            nn.init.trunc_normal_(module.weight, mean=0.0, std=std, a=-cutoff_value, b=cutoff_value)
-        else:
-            nn.init.normal_(module.weight, mean=0.0, std=std)
-    elif config.init_fn == InitFnType.mitchell:
-        std = std_factor / math.sqrt(d)
-        if layer_id is not None:
-            std = std / math.sqrt(2 * (layer_id + 1))
-        nn.init.trunc_normal_(module.weight, mean=0.0, std=std, a=-3 * std, b=3 * std)
-    elif config.init_fn == InitFnType.kaiming_normal:
-        nn.init.kaiming_normal_(module.weight, nonlinearity="relu")
-    elif config.init_fn == InitFnType.fan_in:
-        std = std_factor / math.sqrt(d)
-        nn.init.normal_(module.weight, mean=0.0, std=std)
-    elif config.init_fn == InitFnType.full_megatron:
-        if type_of_module is None:
-            raise RuntimeError(f"When using the {InitFnType.full_megatron} init, every module must have a type.")
-
-        cutoff_factor = config.init_cutoff_factor
-        if cutoff_factor is None:
-            cutoff_factor = 3
-
-        if type_of_module == ModuleType.in_module:
-            # for att_proj (same as QKV), ff_proj
-            std = config.init_std
-        elif type_of_module == ModuleType.out_module:
-            # for attn_out, ff_out
-            std = config.init_std / math.sqrt(2.0 * config.n_layers)
-        elif type_of_module == ModuleType.emb:
-            # positional embeddings (wpe)
-            # token embeddings (wte)
-            std = config.init_std
-        elif type_of_module == ModuleType.final_out:
-            # final output (ff_out)
-            std = config.d_model**-0.5
-        else:
-            raise RuntimeError(f"Unknown module type '{type_of_module}'")
-        nn.init.trunc_normal_(
-            module.weight,
-            mean=0.0,
-            std=std,
-            a=-cutoff_factor * std,
-            b=cutoff_factor * std,
-        )
-    else:
-        raise NotImplementedError(config.init_fn)
-
-    if isinstance(module, nn.Linear):
-        if module.bias is not None:
-            nn.init.zeros_(module.bias)
-
-        if config.init_fn == InitFnType.normal and getattr(module, "_is_residual", False):
-            with torch.no_grad():
-                module.weight.div_(math.sqrt(2 * config.n_layers))