Add configuration resolution and utility functions

simplexity/persistence/utils.py

@@ -0,0 +1,77 @@
+import re
+
+
+def parse_checkpoint_step(path: str) -> int | None:
+    """Extract training step number from checkpoint path.
+
+    Handles multiple formats:
+    - step_12345.pt / step-12345.pt
+    - 12345/model.pt
+    - model_weights/step_00012345.pt
+
+    Args:
+        path: File path or S3 key containing checkpoint
+
+    Returns:
+        Step number if found, None otherwise
+
+    Examples:
+        >>> parse_checkpoint_step("model_weights/step_12345.pt")
+        12345
+        >>> parse_checkpoint_step("checkpoints/12345/model.pt")
+        12345
+        >>> parse_checkpoint_step("step-00500.pt")
+        500
+    """
+    m = re.search(r"step[_-]?(\d+)\.pt$", path)
+    if m:
+        return int(m.group(1))
+
+    parts = path.split("/")
+    if parts and parts[-1] == "model.pt" and len(parts) >= 2:
+        try:
+            return int(parts[-2])
+        except ValueError:
+            pass
+
+    return None
+
+
+def compute_step_width(max_steps: int) -> int:
+    """Compute zero-padding width for step numbers.
+
+    Ensures lexicographic sorting matches chronological order.
+
+    Args:
+        max_steps: Maximum number of training steps
+
+    Returns:
+        Number of digits to use for zero-padding
+
+    Examples:
+        >>> compute_step_width(999)
+        3
+        >>> compute_step_width(100000)
+        6
+    """
+    return len(str(max_steps))
+
+
+def format_step_number(step: int, max_steps: int) -> str:
+    """Format step number with appropriate zero-padding.
+
+    Args:
+        step: Current training step
+        max_steps: Maximum number of training steps
+
+    Returns:
+        Zero-padded step string
+
+    Examples:
+        >>> format_step_number(42, max_steps=100000)
+        '000042'
+        >>> format_step_number(999, max_steps=999)
+        '999'
+    """
+    width = compute_step_width(max_steps)
+    return f"{step:0{width}d}"

simplexity/utils/config_resolution.py

@@ -0,0 +1,67 @@
+def compute_generator_sequence_length(model_n_ctx: int, use_bos: bool) -> int:
+    """Compute the generator's sequence length from model context length and BOS usage.
+
+    The relationship is: model_n_ctx = generator_seq_len - 1 + BOS
+
+    Solving for generator_seq_len: generator_seq_len = model_n_ctx + 1 - BOS
+
+    Args:
+        model_n_ctx: The model's context length (number of input positions it processes)
+        use_bos: Whether a beginning-of-sequence token is prepended during data generation
+
+    Returns:
+        The sequence length to configure for the data generator
+
+    Examples:
+        >>> compute_generator_sequence_length(model_n_ctx=512, use_bos=True)
+        512
+        >>> compute_generator_sequence_length(model_n_ctx=512, use_bos=False)
+        513
+    """
+    return model_n_ctx + 1 - int(use_bos)
+
+
+def compute_model_context_length(generator_seq_len: int, use_bos: bool) -> int:
+    """Compute the model's context length from generator sequence length and BOS usage.
+
+    The relationship is: model_n_ctx = generator_seq_len - 1 + BOS
+
+    Args:
+        generator_seq_len: The sequence length configured for the data generator
+        use_bos: Whether a beginning-of-sequence token is prepended during data generation
+
+    Returns:
+        The context length for the model (number of input positions it will process)
+
+    Examples:
+        >>> compute_model_context_length(generator_seq_len=512, use_bos=True)
+        512
+        >>> compute_model_context_length(generator_seq_len=513, use_bos=False)
+        512
+    """
+    return generator_seq_len - 1 + int(use_bos)
+
+
+def compute_model_vocab_size(generator_vocab_size: int, use_bos: bool, use_eos: bool) -> int:
+    """Compute the model's vocabulary size from generator vocab and special tokens.
+
+    When BOS or EOS tokens are used during data generation, they are added to the vocabulary,
+    increasing the total vocab size the model needs to handle.
+
+    Args:
+        generator_vocab_size: The vocabulary size of the data generator
+        use_bos: Whether a beginning-of-sequence token is used during data generation
+        use_eos: Whether an end-of-sequence token is used during data generation
+
+    Returns:
+        The vocabulary size the model should be configured with
+
+    Examples:
+        >>> compute_model_vocab_size(generator_vocab_size=100, use_bos=True, use_eos=False)
+        101
+        >>> compute_model_vocab_size(generator_vocab_size=100, use_bos=True, use_eos=True)
+        102
+        >>> compute_model_vocab_size(generator_vocab_size=100, use_bos=False, use_eos=False)
+        100
+    """
+    return generator_vocab_size + int(use_bos) + int(use_eos)

simplexity/utils/jnp.py

@@ -3,6 +3,45 @@
 import jax.numpy as jnp
 
 
+def resolve_jax_device(device_spec: str | None = "auto") -> jax.Device:
+    """Resolve device specification to actual JAX device.
+
+    Args:
+        device_spec: One of "auto", "gpu", "cuda", "cpu", or None (treated as "auto")
+
+    Returns:
+        JAX device object
+
+    Examples:
+        >>> resolve_jax_device("auto")  # On GPU machine
+        GpuDevice(id=0, ...)
+        >>> resolve_jax_device("cpu")
+        CpuDevice(id=0)
+    """
+    if device_spec is None or device_spec == "auto":
+        try:
+            devices = jax.devices("gpu")
+            if devices:
+                return devices[0]
+        except RuntimeError:
+            pass
+        return jax.devices("cpu")[0]
+
+    if device_spec in ("gpu", "cuda"):
+        try:
+            devices = jax.devices("gpu")
+            if devices:
+                return devices[0]
+        except RuntimeError:
+            pass
+        raise RuntimeError("GPU requested but no GPU devices available")
+
+    if device_spec == "cpu":
+        return jax.devices("cpu")[0]
+
+    raise ValueError(f"Unknown device specification: {device_spec}")
+
+
 @eqx.filter_jit
 def entropy(probs: jax.Array, log: bool = False) -> jax.Array:
     """Compute the entropy of a log probability distribution."""

simplexity/utils/pytorch_utils.py

@@ -78,3 +78,46 @@ def torch_to_jax(torch_tensor: torch.Tensor) -> jax.Array:
         numpy_array = torch_tensor.detach().cpu().numpy()
         jax_array = jnp.array(numpy_array)
         return jax_array
+
+
+def resolve_device(device_spec: str | None = "auto") -> str:
+    """Resolve device specification to actual PyTorch device string.
+
+    Args:
+        device_spec: One of "auto", "cuda", "mps", "cpu", or None (treated as "auto")
+
+    Returns:
+        Resolved device string: "cuda", "mps", or "cpu"
+
+    Raises:
+        ValueError: If device_spec is not a recognized device type
+        RuntimeError: If a specific device is requested but unavailable
+
+    Examples:
+        >>> resolve_device("auto")  # On CUDA machine
+        'cuda'
+        >>> resolve_device("cpu")
+        'cpu'
+    """
+    if device_spec is None or device_spec == "auto":
+        if torch.cuda.is_available():
+            return "cuda"
+        elif torch.backends.mps.is_available():
+            return "mps"
+        else:
+            return "cpu"
+
+    if device_spec == "cuda":
+        if not torch.cuda.is_available():
+            raise RuntimeError("CUDA requested but CUDA is not available")
+        return "cuda"
+
+    if device_spec == "mps":
+        if not torch.backends.mps.is_available():
+            raise RuntimeError("MPS requested but MPS is not available")
+        return "mps"
+
+    if device_spec == "cpu":
+        return "cpu"
+
+    raise ValueError(f"Unknown device specification: {device_spec}")