Initial K3 snapshot: 0.5B KDA/MLA/MoE train path

Standalone tree split from LLMRL/projects/kda. Includes Triton dt_bias
backward fix, train_k3 --preset 0.5b, SFT, Docker runtime, and tests.

kda/training/schedule.py

@@ -0,0 +1,47 @@
+"""LR scale and token-horizon helpers for train_k3 / train_sft."""
+from __future__ import annotations
+
+import math
+
+
+def lr_scale(
+    opt_step: int,
+    warmup: int,
+    total_opt: int,
+    min_ratio: float = 0.1,
+) -> float:
+    """Linear warmup (optimizer steps) then cosine down to ``min_ratio``.
+
+    ``opt_step`` is 0-indexed at the optimizer update that is about to run.
+    """
+    if warmup > 0 and opt_step < warmup:
+        return (opt_step + 1) / warmup
+    denom = max(total_opt - warmup - 1, 1)
+    progress = min(max(opt_step - warmup, 0) / denom, 1.0)
+    cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
+    return min_ratio + (1.0 - min_ratio) * cosine
+
+
+def tokens_per_micro(batch: int, seq_len: int) -> int:
+    return batch * seq_len
+
+
+def total_opt_steps(
+    *,
+    max_tokens: int | None,
+    max_micro: int | None,
+    batch: int,
+    seq_len: int,
+    grad_acc: int,
+) -> int:
+    """Optimizer-step horizon used by cosine. At least 1."""
+    acc = max(grad_acc, 1)
+    candidates: list[int] = []
+    if max_tokens is not None and max_tokens > 0:
+        tpm = max(tokens_per_micro(batch, seq_len), 1)
+        candidates.append(math.ceil(max_tokens / (tpm * acc)))
+    if max_micro is not None and max_micro > 0:
+        candidates.append(math.ceil(max_micro / acc))
+    if not candidates:
+        return 1
+    return max(min(candidates), 1)