feat: torchJD working version

ValerianRey · ValerianRey · commit 37fe9eb0718e · 2025-12-22T19:16:44.000Z
diff --git a/.vscode/launch.json b/.vscode/launch.json
@@ -18,7 +18,7 @@
                 "epochs=50000",
                 "eval_interval=5000",
                 "lr=1e-4",
-                "global_batch_size=768",
+                "global_batch_size=1",
                 "puzzle_emb_lr=1e-4",
                 "weight_decay=1.0",
                 "puzzle_emb_weight_decay=1.0",
diff --git a/pyproject.toml b/pyproject.toml
@@ -26,4 +26,5 @@ dependencies = [
 	"numba",
 	"triton",
 	"pre-commit",
+	"torchjd>=0.8.0",
 ]
diff --git a/src/recursion/pretrain.py b/src/recursion/pretrain.py
@@ -16,6 +16,7 @@
 from adam_atan2_pytorch import AdamAtan2
 from omegaconf import DictConfig
 from torch import nn
+from torch.nn.functional import cosine_similarity
 from torch.utils.data import DataLoader
 from torchjd.aggregation import UPGrad
 from torchjd.autojac._transform import Accumulate, Aggregate, OrderedSet
@@ -26,6 +27,35 @@
 from recursion.puzzle_dataset import PuzzleDataset, PuzzleDatasetConfig, PuzzleDatasetMetadata
 from recursion.utils.functions import get_model_source_path, load_model_class
 
+global_step = 0
+
+
+def print_gramian(_, inputs, __):
+    if not dist.is_initialized() or dist.get_rank() == 0:
+        # print(inputs[0])
+        diag = torch.diag(inputs[0]).sqrt()
+        outer = diag.unsqueeze(0) * diag.unsqueeze(1)
+        scaled_gramian = inputs[0].clone() / outer
+
+        wandb.log(
+            {
+                "gramian_min": inputs[0].min(),
+                "gramian_mean": inputs[0].mean(),
+                "gramian_median": inputs[0].median(),
+                "min_gramian_scaled": scaled_gramian.min(),
+            },
+            step=global_step,
+        )
+
+
+def log_gd_similarity(_, inputs: tuple[torch.Tensor, ...], aggregation: torch.Tensor) -> None:
+    """Prints the cosine similarity between the aggregation and the average gradient."""
+    if not dist.is_initialized() or dist.get_rank() == 0:
+        matrix = inputs[0]
+        gd_output = matrix.mean(dim=0)
+        similarity = cosine_similarity(aggregation, gd_output, dim=0)
+        wandb.log({"gd_similarity": similarity.item()}, step=global_step)
+
 
 class LossConfig(pydantic.BaseModel):
     model_config = pydantic.ConfigDict(extra="allow")
@@ -320,7 +350,7 @@ def create_evaluators(config: PretrainConfig, eval_metadata: PuzzleDatasetMetada
     return evaluators
 
 
-UPDATE_EVERY = 2
+UPDATE_EVERY = 8
 
 
 def train_batch(
@@ -331,7 +361,9 @@ def train_batch(
     rank: int,
     world_size: int,
 ):
+    global global_step
     train_state.step += 1
+    global_step = train_state.step
     if train_state.step > train_state.total_steps:  # At most train_total_steps
         return
 
@@ -348,8 +380,8 @@ def train_batch(
         carry=train_state.carry, batch=batch, return_keys=[]
     )
 
-    current_step = train_state.carry.steps[0].item()  # Something between 0 and 15
-    if (current_step + 1) % UPDATE_EVERY:
+    current_step = train_state.carry.steps[0].item()  # Something between 1 and 16
+    if current_step > 0 and (current_step % UPDATE_EVERY == 0):
         memories = train_state.carry.inner_carry.memories
         memories_wrt = train_state.carry.inner_carry.memories_wrt
         rec_model = train_state.model.model.inner.L_level
@@ -375,6 +407,8 @@ def train_batch(
             param: torch.stack(gradients, dim=0) for param, gradients in param_to_gradients.items()
         }
         aggregator = UPGrad()
+        aggregator.weighting.weighting.weighting.register_forward_hook(print_gramian)
+        aggregator.register_forward_hook(log_gd_similarity)
         transform = Accumulate() << Aggregate(aggregator, OrderedSet(list(rec_model.parameters())))
         transform(param_to_jacobian)  # This stores the aggregated Jacobian in the .grad fields
 

Original file line number	Diff line number	Diff line change
`@@ -26,4 +26,5 @@ dependencies = [`
`26`	`26`	`"numba",`
`27`	`27`	`"triton",`
`28`	`28`	`"pre-commit",`
	`29`	`+ "torchjd>=0.8.0",`
`29`	`30`	`]`