Fix quantize-lm on macOS (#21), gradient checkpointing freezing blocks, RoPE grader accepting no-ops

[Tar-ive]

Fixes #21. Also fixes gradient checkpointing silently leaving every checkpointed block untrained, and tightens the RoPE grader so a no-op implementation no longer passes.

1. quantize-lm fails on macOS / torch ≥ 2.6 (#21)

Reproduced on an M-series Mac with torch 2.9.0 (engine defaults to 'none', supported: ['qnnpack']), in two stages exactly as reported:

• quantize_dynamic → RuntimeError: Didn't find engine for operation quantized::linear_prepack NoQEngine
• after setting the engine, torch.load(...) → UnpicklingError ... Unsupported global: GLOBAL torch.ScriptObject. A dynamically quantized LSTM state_dict stores packed params as torch.ScriptObject, which weights_only=True (the default since 2.6) rejects.

Fix (question and solution notebooks share this scaffold, so both are updated):

• Select a supported quantized engine when it is 'none'. This is the same guard the grader already has in quantize_lm.py.
• Save the float weights and quantize again after loading. Dynamic quantization is deterministic, so the reloaded model is identical, and the load stays on the safe weights_only=True path. The alternative, weights_only=False or allow-listing torch.ScriptObject, would teach an unsafe loading habit.
• from torch.ao.quantization import quantize_dynamic (torch.quantization is the deprecated alias).

Separate bug in the same notebook: the model ends in nn.Softmax and was trained with nn.CrossEntropyLoss, which applies a second softmax. That pinned the loss at ≈ ln(50); the committed output showed 3.9118 → 3.9097 over 5 epochs. The grader requires probabilities out of the model, so the fix keeps the softmax and trains with NLLLoss on log(probs). The loss now goes 3.9032 → 3.8048.

The solution also gains the comparison the problem statement asks for ("evaluate the quantized model's performance compared to the original model"): on this machine, weights go 0.97 MB → 0.27 MB with 100% top-1 agreement. Outputs produced by the old code were cleared.

2. Gradient checkpointing: checkpointed blocks never train

autograd only records a custom Function, and only calls its backward, if at least one tensor argument requires grad. The block's parameters are closed over by fn, not passed as arguments. So when the input is plain data, which is the notebook's own training test and the first layer of any real network, backward never ran. Only head trained, all 96 block parameters had grad is None, and the test still passed because the head alone lowered the loss.

• Solution: checkpoint() passes a dummy requires_grad=True tensor when no argument requires grad. backward skips outputs that need no grad (e.g. a frozen fn).
• Training test: asserts every parameter receives a gradient with plain input. The previous solution fails it with 96 parameters got no gradient.
• Question notebook: the hint and TODOs no longer steer to torch.autograd.grad w.r.t. the inputs, which never reaches fn's parameters. The validation now treats a missing gradient as a failure instead of skipping it, matching the solution.
• Grader: check_module_parameter_gradients used to raise Skip when the wrapped module got no gradient. The runner counts that as a pass, so the broken formulation was accepted. It now fails. A new check, check_plain_input_still_trains_parameters, covers the no-grad input case. The hints were corrected and regenerated into problem.toml / problems.json (generate.py hints + json). They were also synced into the MCP server's hints.generated.ts, which no script regenerates today.

3. RoPE grader accepted a no-op

check("rotary-positional-embedding", rotate_half, apply) passed with apply = lambda q, k, cos, sin: (q, k). The zero-angle check expects exactly that, and a no-op trivially preserves norms. It also passed when only q was rotated. The new check_quarter_turn checks that at cos=0, sin=1 the output equals the solver's own rotate_half(q) and rotate_half(k), so any consistent pairing convention still passes.

Tests

All run on torch 2.9.0, macOS arm64, CPU:

• python/test_grader.py: 14/14. This includes 5 new wrong-answer cases (RoPE no-op, RoPE query-only, checkpoint without dummy input, checkpoint with autograd.grad on inputs only, plus the correct versions), so these checks can't regress to accepting broken code.
• scripts/validate_specs.py rotary-positional-embedding gradient-checkpointing quantize-lm: all 3 repo solutions pass their updated specs.
• Both fixed solution notebooks run end to end. The quantize-lm question scaffold also runs end to end with the solution's model.
• generate.py check: 68 manifests, 0 problems. The MCP server typechecks (tsc --noEmit).

Notes / follow-ups (not in this PR)

• torch.ao.quantization emits a deprecation warning announcing its removal. Longer term this exercise should move to torchao (quantize_ with an int8 dynamic config).
• The PyPI package is still 0.1.0, so the grader changes reach users only after a release.
• Issue In h3, we should pass batch_first=True to TransformerEncoderLayer #5 (transformer batch_first) looks already fixed on main (60d933e replaced nn.TransformerEncoderLayer with a hand-written batch-first layer; the current solution reaches ~98.5%) and could be closed.

[vercel]

@Tar-ive is attempting to deploy a commit to the exorust's projects Team on Vercel.

A member of the Team first needs to authorize it.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB