freesolo-chalk 0.2.0

Custom Triton/CUDA kernels that complement Liger Kernel for LLM post-training

Chalk

Custom Triton/CUDA kernels that complement Liger Kernel.

pip install freesolo-chalk

Liger fuses the cross-entropy, activation, and RMSNorm paths. Chalk fills the gaps that matter for Freesolo's Flash post-training stack — fused GEMMs, the LoRA-delta matmuls, the QKV norm+RoPE epilogue, embedding gather, and FP8 frozen-base GEMMs — each behind a documented, benchmarked, opt-in entry point.

Chalk's repository layout and conventions intentionally mirror Liger Kernel.

Layout

src/chalk/
  ops/           # raw Triton/CUDA kernels + autograd.Function wrappers
  transformers/  # model-level installers that monkeypatch kernels into HF modules
  utils.py       # device detection helpers
test/            # correctness + gating tests (mirrors test/ops, test/transformers)
benchmark/       # speed + fp32-correctness A/B harness

Design principles

• Worker-side kernel library. Like Liger, chalk depends on torch + triton — it is meant to be installed where kernels actually run (the GPU worker), so consumers should depend on it from a gpu extra rather than their base install. Importing the top-level chalk package is still cheap (kernels lazy-load), so probing chalk.utils.infer_device() never forces a heavy import.
• Complements, not replaces, Liger. Liger fuses CE / activation / RMSNorm; chalk fuses the GEMMs, LoRA delta, QKV epilogue, embedding, and FP8 base.
• Safe fallback. Every installer is arch-gated, runs a numeric self-test on install, patches only frozen nn.Linear layers (never trainable / PEFT-wrapped layers), and silently falls back to the eager / Liger path on any import / compile / self-test failure.
• Opt-in & evidence-based. Kernels are off unless explicitly enabled, and every kept kernel has end-to-end loss-curve evidence — not just a microbenchmark.

Development

pip install -e '.[dev]'
make checkstyle   # ruff check + format
make test         # pytest with coverage
make run-benchmarks

Status

Intentionally minimal to start — kernels are landed one at a time under chalk/ops + chalk/transformers, each with correctness tests and benchmark evidence.

License

BSD-2-Clause. See LICENSE and NOTICE.