Repository: allenai/OLMo
File: olmo/config.py (Python)

Code before:
    a model instance is used per node (as determined by ``get_world_size() // get_local_world_size()``).
    PyTorch's default HSDP behavior matches this default behavior.
    """


@dataclass
class SingleGPUConfig(BaseConfig):
    device: str = "auto"
    """
    Device to run single-device training.
    """

    def get_device(self):
        if self.device == "auto":
            if torch.backends.mps.is_available():
                return torch.device("mps")
            elif torch.cuda.is_available():
                return torch.device("cuda")
            else:
                return torch.device("cpu")
        elif self.device == "mps" and not torch.backends.mps.is_available():
            print("MPS not available.")
            return torch.device("cpu")
        elif self.device == "cuda" and not torch.cuda.is_available():
            print("CUDA not available.")
            return torch.device("cpu")
        else:
            return torch.device(self.device)


class CheckpointType(StrEnum):
    sharded = "sharded"
    unsharded = "unsharded"
    sharded_ephemeral = "sharded_ephemeral"


class ShardedCheckpointerType(StrEnum):
    torch_new = "torch_new"
    torch_legacy = "torch_legacy"
    local = "local"
    olmo_core = "olmo_core"


class ActivationCheckpointingStrategy(StrEnum):
    whole_layer = "whole_layer"
    """
    Checkpoint every transformer layer.
    """

    one_in_two = "one_in_two"
    """

Review comment (bug):
If the device is explicitly set to something that won't work, we should just fail instead of silently fixing it. Throw an exception in this case, like `OLMoConfigurationError`.

We should _never_ just print something. Sometimes we log something (through the logging mechanism). If we need things to fail, we throw exceptions. Also, printing something will go to stdout, where it does not belong in the case of failures.

Code after:

    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: Optional[FSDPPrecision] = FSDPPrecision.pure

    hybrid_sharding_num_model_replicas: Optional[int] = None
    """
    The number of model instances, when using a hybrid sharding strategy.
    If not ``None``, this must divide the total number of nodes. If ``None``, the default,
    a model instance is used per node (as determined by ``get_world_size() // get_local_world_size()``).
    PyTorch's default HSDP behavior matches this default behavior.
    """


@dataclass
class SingleGPUConfig(BaseConfig):
    device: str = "auto"
    """
    Device to run single-device training.
    """

    def get_device(self):
        if self.device == "auto":
            if torch.backends.mps.is_available():
                return torch.device("mps")
            elif torch.cuda.is_available():
                return torch.device("cuda")
            else:
                return torch.device("cpu")
        elif self.device == "mps" and not torch.backends.mps.is_available():
            raise OLMoConfigurationError("MPS not available.")