fix: patch torch._pytree.register_constant + tighten torchao/transformers pins

- Add no-op shim for torch.utils._pytree.register_constant in train.py
before any torchao/unsloth/transformers import. Fixes AttributeError
when torchao>=0.9 is installed with PyTorch 2.6.0.

- Dockerfile: move torchao==0.6.1 force-reinstall to LAST pip step so
nothing can override it after the fact.

- Dockerfile: tighten transformers pin to <4.48 (4.48+ unconditionally
imports quantizer_torchao at module init, breaking torchao 0.6.1).

- Dockerfile: remove --upgrade from 'pip install -r requirements.txt'
to prevent silent package version bumps.

Dockerfile

@@ -12,19 +12,23 @@ WORKDIR /app
 # Upgrade pip
 RUN pip install --no-cache-dir --upgrade pip
 
-# Install unsloth from GitHub (it will drag in a bleeding-edge torchao)
+# Install unsloth from GitHub (drags in bleeding-edge torchao)
 RUN pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"
 
-# CRITICAL: Force-downgrade torchao to last known-good version AFTER unsloth
-# torchao>=0.7 requires torch.utils._pytree.register_constant which does not
-# exist in any stable PyTorch release yet.
-RUN pip install --no-cache-dir --force-reinstall "torchao==0.6.1"
-
-# Pin trl and transformers to stable versions
-RUN pip install --no-cache-dir "trl>=0.15,<0.17" "transformers>=4.47,<4.51"
+# Pin trl and transformers BEFORE locking torchao.
+# transformers>=4.48 unconditionally does `from .quantizer_torchao import ...`
+# at module init, which breaks if torchao is the wrong version.
+# Stay on <4.48 where that import is guarded.
+RUN pip install --no-cache-dir "trl>=0.15,<0.17" "transformers>=4.47,<4.48"
 
 COPY --chown=user ./requirements.txt requirements.txt
-RUN pip install --no-cache-dir --upgrade -r requirements.txt
+# No --upgrade here; we don't want pip silently re-upgrading pinned packages.
+RUN pip install --no-cache-dir -r requirements.txt
+
+# CRITICAL: Force-downgrade torchao LAST (after every other install).
+# torchao>=0.7 requires torch.utils._pytree.register_constant which does not
+# exist in PyTorch 2.6.0.  This must be the final pip step. (v2)
+RUN pip install --no-cache-dir --force-reinstall "torchao==0.6.1"
 
 # Copy all the code
 COPY --chown=user . /app

trainer/train.py

@@ -1,6 +1,16 @@
 import sys
 import os
 
+# ── Compatibility shim ──────────────────────────────────────────────────────
+# torchao>=0.9 decorates enums with @register_as_pytree_constant which calls
+# torch.utils._pytree.register_constant – a function that only exists in
+# PyTorch 2.7+.  We're running on 2.6.0, so patch it in before any import
+# of torchao / transformers / unsloth triggers the missing attribute error.
+import torch.utils._pytree as _pytree
+if not hasattr(_pytree, "register_constant"):
+    _pytree.register_constant = lambda cls: cls  # no-op shim
+# ────────────────────────────────────────────────────────────────────────────
+
 os.environ["WANDB_API_KEY"] = "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"
 os.environ["WANDB_PROJECT"] = "codeforge-grpo"