ML Systems Layer — Scaling Studies (MLSys1)

.flake8

@@ -1,5 +1,5 @@
 [flake8]
-max-line-length = 88
+max-line-length = 100
 extend-ignore = E203, W503
 
 exclude =

1_mini_clip/src/mini_clip/loss/contrastive.py

@@ -5,31 +5,62 @@
 
 class ClipLoss(nn.Module):
     """
-    Contrastive Loss (InfoNCE) for CLIP training.
-    Computes the symmetric cross-entropy loss between image and text features.
+    Backwards-compatible CLIP contrastive loss.
+
+    Accepts logit_scale in either form:
+      1) raw/log-space (typical): model.logit_scale (around ~2-5)
+      2) already-exp'd scale: exp(model.logit_scale) (around ~1-100)
+
+    Heuristic:
+      - if logit_scale <= ~20, treat as raw and exp()
+      - else treat as already-exp'd
     """
 
-    def __init__(self):
+    def __init__(self, max_scale: float = 100.0):
         super().__init__()
+        self.max_scale = float(max_scale)
+        self._max_log = float(
+            torch.log(torch.tensor(self.max_scale)).item()
+        )  # ln(max_scale)
+
+    def _as_scale(self, logit_scale: torch.Tensor) -> torch.Tensor:
+        s = logit_scale.float()
+
+        # If it's already a scale (e.g., 100), we should NOT exp again.
+        # If it's raw log-scale (~2-5), we SHOULD exp.
+        # 20 is a safe separator: exp(20) is enormous,
+        # raw logit_scale won't be that high in sane CLIP.
+        if s.item() <= 20.0:
+            s = s.clamp(max=self._max_log).exp()
+        else:
+            s = s.clamp(max=self.max_scale)
 
-    def forward(self, image_features, text_features, logit_scale):
-        """
-        Args:
-            image_features: [batch_size, dim] normalized image features
-            text_features: [batch_size, dim] normalized text features
-            logit_scale: scalar logit scale (exp(model.logit_scale))
-        """
+        return s
+
+    def forward(
+        self,
+        image_features: torch.Tensor,
+        text_features: torch.Tensor,
+        logit_scale: torch.Tensor,
+    ):
         device = image_features.device
-        logits_per_image = logit_scale.exp() * image_features @ text_features.T
+
+        # AMP stability: do logits + CE in fp32
+        img = image_features.float()
+        txt = text_features.float()
+        scale = self._as_scale(logit_scale)
+
+        logits_per_image = scale * (img @ txt.T)
         logits_per_text = logits_per_image.T
 
-        batch_size = logits_per_image.shape[0]
-        labels = torch.arange(batch_size, device=device, dtype=torch.long)
+        bsz = logits_per_image.size(0)
+        labels = torch.arange(bsz, device=device, dtype=torch.long)
 
         loss_img = F.cross_entropy(logits_per_image, labels)
         loss_txt = F.cross_entropy(logits_per_text, labels)
 
-        total_loss = (loss_img + loss_txt) / 2
+        total_loss = 0.5 * (loss_img + loss_txt)
+
         return {
             "loss": total_loss,
             "loss_img": loss_img,

2_ml_systems_mini_clip/src/ml_systems/README.md

@@ -0,0 +1 @@
+# Currently implemented inside orchestration/ml_systems