freesolo-chalk 0.5.2.dev0

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)

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, and silently falls back to the eager / Liger path on any import / compile / self-test failure. All installers but one patch only frozen nn.Linear layers (never trainable / PEFT-wrapped layers). The exception is the LoRA-delta installer (install_fused_lora_delta), the one kernel that intentionally accelerates the trainable PEFT LoRA path: it monkeypatches PEFT's dense lora.Linear.forward to route the adapter delta through the fused kernel. Its self-test checks the fused forward and the x / A / B gradients against torch autograd (bf16 tolerance) before patching, and at runtime it falls back to stock PEFT per call for everything the kernel does not replicate exactly (merged / disabled adapters, mixed-adapter batches, LoRA variants, dropout, lora_bias, quantized / non-dense bases, and any non-CUDA tensor) — so a failed self-test or any unsupported layer leaves training numerics on the stock PEFT path.
• 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

Status

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

License

BSD-2-Clause. See LICENSE and NOTICE.