use 1-sqrt warmdown shape for LR schedule

[spjosyula]

Replace linear warmdown with 1-sqrt cooldown shape (LR = 1 - sqrt(x)) in
base_train.py and chat_sft.py. Left chat_rl.py unchanged since its schedule is pure decay with no stable phase.
The warmdown shape has never been changed since it was inherited from
modded-nanogpt: the ratio was swept (0.2 → 0.4 → 0.5) but the linear
curve itself was kept as-is. 1-sqrt drops LR faster early in warmdown
then flattens, which should reduce gradient noise sooner while keeping
more steps at low-but-nonzero LR for convergence.
Needs a d12 run to validate.

[svlandeg]

Have you done any experiments to validate the change?

[spjosyula]

Have you done any experiments to validate the change?

No experiments yet. As noted in the description, this needs a d12 run.

Why I think it's worth one:

• The LR multiplier applies to all param groups uniformly, so the shape change affects every parameter across
  both AdamW and Muon
• Warmdown is 50% of training: that's a large surface for the curve shape to matter
• 1-sqrt outperformed linear in WSD schedules (Hagele et al., NeurIPS 2024 - already cited in the code at L346)

The gap is that those results were with AdamW and if the shape shifts the optimal warmdown ratio, that may need
re-sweeping too.