pyproject.toml

[project]
name = "academic-pretraining"
version = "1.1.0"
description = "Codebase for the paper '$100K or 100 Days: Trade-offs when Pre-Training with Academic Resources'"
readme = "README.md"
requires-python = ">=3.12"
license = { text = "MIT" }
authors = [{ name = "Apoorv Khandelwal", email = "mail@apoorvkh.com" }]
dependencies = [
  # training libraries
  "accelerate==1.2.1",
  "deepspeed==0.16.2",
  "tabulate==0.9.0",      # for counting flops in pytorch
  "torch==2.5.1",
  "torchvision==0.20.1",
  "transformers==4.47.1",
  "triton==3.1.0",
  # experiment execution
  "ai2-tango==1.3.2",
  "torchrunx==0.2.3",
  "submitit==1.5.1",
  "tyro==0.9.4",
  # plotting
  "marimo==0.10.6",
  "polars[style]==1.17.1",
  "altair[save]==5.5.0",
  # linting, formatting, type checking
  "ruff==0.8.4",
  "pyright==1.1.391",
]

[project.optional-dependencies]
# custom kernels for mamba
mamba = [
  "causal-conv1d @ https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.5.0.post8/causal_conv1d-1.5.0.post8+cu12torch2.5cxx11abiFALSE-cp312-cp312-linux_x86_64.whl",
  "mamba-ssm @ https://github.com/state-spaces/mamba/releases/download/v2.2.4/mamba_ssm-2.2.4+cu12torch2.5cxx11abiFALSE-cp312-cp312-linux_x86_64.whl",
]

[tool.pixi.project]
channels = ["nvidia/label/cuda-12.4.0", "conda-forge"]
platforms = ["linux-64"]

[tool.pixi.dependencies]
gcc = "11.4.0"
gxx = "11.4.0"
python = "3.12.8"
pip = "24.3.1"
uv = "0.5.11"
cuda = "12.4.0"

[tool.pixi.environments]
default = { features = ["mamba"], solve-group = "default" }

[tool.pixi.activation]
scripts = [".env"] # auto `source .env` in `pixi shell`

[tool.ruff]
line-length = 120
src = ["src", "experiments", "scripts"]
target-version = "py312"

[tool.ruff.lint]
select = ["E", "F", "I"]
ignore = ["E402"]

[tool.ruff.lint.per-file-ignores]
"__init__.py" = ["F401"]
"scripts/plotting/*.py" = [
  "E712", # Allow `== bool` in Polars
]

[tool.pyright]
include = ["src", "experiments", "scripts"]
exclude = ["scripts/plotting"]
pythonVersion = "3.12"
reportRedeclaration = "none"