From fbda1c83a48289757db4333924f0c2677a17235a Mon Sep 17 00:00:00 2001
From: krmayankb <krmayankb@gmail.com>
Date: Fri, 21 Jul 2023 11:56:19 -0700
Subject: [PATCH 01/18] mrl integration

---
 scripts/launch_this.slurm | 16 ++++++++++++++++
 scripts/mrl_clip.sh       | 22 ++++++++++++++++++++++
 src/open_clip/factory.py  | 19 +++++++++++++++++--
 src/open_clip/loss.py     | 35 +++++++++++++++++++++++++++++++++++
 src/open_clip/model.py    | 12 ++++++++++--
 src/training/main.py      |  1 +
 src/training/params.py    | 14 ++++++++++++++
 7 files changed, 115 insertions(+), 4 deletions(-)
 create mode 100644 scripts/launch_this.slurm
 create mode 100644 scripts/mrl_clip.sh

diff --git a/scripts/launch_this.slurm b/scripts/launch_this.slurm
new file mode 100644
index 0000000..16e8be1
--- /dev/null
+++ b/scripts/launch_this.slurm
@@ -0,0 +1,16 @@
+#!/bin/bash
+
+#SBATCH --account=krishna
+#SBATCH --partition=gpu-a100
+#SBATCH --job-name=mrl
+#SBATCH --error=./jobs/mrl_clip/job.err.%j  
+#SBATCH --output=./jobs/mrl_clip/job.run.%j
+#SBATCH --time=10-00:00:00
+#SBATCH --cpus-per-task=6
+#SBATCH --gpus=2
+#SBATCH --mem-per-gpu=80G
+#SBATCH --nodes=1
+#SBATCH --mail-user=krmayank@uw.edu
+#SBATCH --mail-type=ALL
+
+bash mrl_clip.sh
\ No newline at end of file
diff --git a/scripts/mrl_clip.sh b/scripts/mrl_clip.sh
new file mode 100644
index 0000000..fe6d59b
--- /dev/null
+++ b/scripts/mrl_clip.sh
@@ -0,0 +1,22 @@
+cd /mmfs1/gscratch/krishna/mayank/clip_clone/open_clip/src
+torchrun --nproc_per_node 2 --master_port 3234 -m training.main \
+    --model "ViT-B-16" \
+    --train-data "/mmfs1/data/yfcc-tmp/cc_3m/train_shards/shard_{000000..003318}.tar" \
+    --imagenet-val "/mmfs1/data/yfcc-tmp/imagenet/val/" \
+    --dataset-type webdataset \
+    --precision amp \
+    --gather-with-grad \
+    --local-loss \
+    --force_mrl_loss \
+    --mrl_loss_weights "1,1,1,1,1,1,1,1" \
+    --batch-size 512 \
+    --accum-freq 1 \
+    --workers 4 \
+    --epochs 40 \
+    --warmup 2000 \
+    --zeroshot-frequency 2 \
+    --seed 0 \
+    # --report-to 'wandb' \
+    # --wandb-project-name "mrl_clip_training" \
+    # --logs "/mmfs1/gscratch/krishna/mayank/clip_clone/open_clip/src/logs/mrl_clip" \
+    # --name "mrl_clip_b512_accum_1_ep40" 
\ No newline at end of file
diff --git a/src/open_clip/factory.py b/src/open_clip/factory.py
index 14011f9..f1fda90 100644
--- a/src/open_clip/factory.py
+++ b/src/open_clip/factory.py
@@ -13,7 +13,7 @@
 from .model import CLIP, CustomTextCLIP, convert_weights_to_lp, convert_to_custom_text_state_dict,\
     resize_pos_embed, get_cast_dtype
 from .coca_model import CoCa
-from .loss import ClipLoss, DistillClipLoss, CoCaLoss
+from .loss import ClipLoss, DistillClipLoss, CoCaLoss, MRLClipLoss
 from .openai import load_openai_model
 from .pretrained import is_pretrained_cfg, get_pretrained_cfg, download_pretrained, list_pretrained_tags_by_model, download_pretrained_from_hf
 from .transform import image_transform, AugmentationCfg
@@ -119,6 +119,7 @@ def create_model(
         cache_dir: Optional[str] = None,
         output_dict: Optional[bool] = None,
         require_pretrained: bool = False,
+        use_mrl: bool = False
 ):
     has_hf_hub_prefix = model_name.startswith(HF_HUB_PREFIX)
     if has_hf_hub_prefix:
@@ -191,7 +192,7 @@ def create_model(
             else:
                 model = CustomTextCLIP(**model_cfg, cast_dtype=cast_dtype)
         else:
-            model = CLIP(**model_cfg, cast_dtype=cast_dtype)
+            model = CLIP(**model_cfg, cast_dtype=cast_dtype, use_mrl=use_mrl) # intialise with mrl based config if use_mrl
 
         pretrained_loaded = False
         if pretrained:
@@ -261,6 +262,16 @@ def create_loss(args):
             world_size=args.world_size,
             use_horovod=args.horovod,
         )
+    elif args.force_mrl_loss: 
+        return MRLClipLoss(
+            local_loss=args.local_loss,
+            gather_with_grad=args.gather_with_grad,
+            cache_labels=True,
+            rank=args.rank,
+            world_size=args.world_size,
+            use_horovod=args.horovod,
+            mrl_loss_weights=args.mrl_loss_weights
+        )
     return ClipLoss(
         local_loss=args.local_loss,
         gather_with_grad=args.gather_with_grad,
@@ -288,6 +299,7 @@ def create_model_and_transforms(
         aug_cfg: Optional[Union[Dict[str, Any], AugmentationCfg]] = None,
         cache_dir: Optional[str] = None,
         output_dict: Optional[bool] = None,
+        use_mrl = False
 ):
     model = create_model(
         model_name,
@@ -303,6 +315,7 @@ def create_model_and_transforms(
         pretrained_hf=pretrained_hf,
         cache_dir=cache_dir,
         output_dict=output_dict,
+        use_mrl=use_mrl
     )
 
     image_mean = image_mean or getattr(model.visual, 'image_mean', None)
@@ -337,6 +350,7 @@ def create_model_from_pretrained(
         image_mean: Optional[Tuple[float, ...]] = None,
         image_std: Optional[Tuple[float, ...]] = None,
         cache_dir: Optional[str] = None,
+        use_mrl = False
 ):
     model = create_model(
         model_name,
@@ -349,6 +363,7 @@ def create_model_from_pretrained(
         force_image_size=force_image_size,
         cache_dir=cache_dir,
         require_pretrained=True,
+        use_mrl=use_mrl
     )
 
     if not return_transform:
diff --git a/src/open_clip/loss.py b/src/open_clip/loss.py
index 4fbf61d..b1e9690 100644
--- a/src/open_clip/loss.py
+++ b/src/open_clip/loss.py
@@ -210,3 +210,38 @@ def forward(
             return {"contrastive_loss": contrastive_loss, "distill_loss": distill_loss}
 
         return contrastive_loss, distill_loss
+
+class MRLClipLoss(ClipLoss): 
+    def __init__(self, 
+                 local_loss=False, 
+                 gather_with_grad=False, 
+                 cache_labels=False, 
+                 rank=0, 
+                 world_size=1, 
+                 use_horovod=False, 
+                 mrl_loss_weights = None):
+        super().__init__(local_loss, 
+                         gather_with_grad, 
+                         cache_labels, 
+                         rank, 
+                         world_size, 
+                         use_horovod)
+        self.mrl_loss_weights = mrl_loss_weights
+    
+    def forward(self, image_features, text_features, logit_scale, output_dict=False):
+        # print("Inside forward of MRL CLIP loss",len(self.mrl_loss_weights), self.mrl_loss_weights )
+
+        assert len(self.mrl_loss_weights) == 8, "list containing loss weights must contain 8 elements"
+        
+        dim_to_consider = [8, 16, 32, 64, 128, 256, 512, 768]
+        total_loss = 0 
+        for idx, dim in enumerate(dim_to_consider): 
+            img = F.normalize( image_features[:,:dim], dim=-1) # slice and normalize 
+            text = F.normalize( text_features[:,:dim], dim=-1) # slice and normalize 
+            loss = super().forward(image_features=img, text_features=text, logit_scale=logit_scale)
+            total_loss += self.mrl_loss_weights[idx] * loss
+        
+        if output_dict:
+            return {"mrl_clip_loss": total_loss}        
+        
+        return total_loss
diff --git a/src/open_clip/model.py b/src/open_clip/model.py
index 4f5e775..7729e14 100644
--- a/src/open_clip/model.py
+++ b/src/open_clip/model.py
@@ -184,9 +184,11 @@ def __init__(
             quick_gelu: bool = False,
             cast_dtype: Optional[torch.dtype] = None,
             output_dict: bool = False,
+            use_mrl: bool = False
     ):
         super().__init__()
         self.output_dict = output_dict
+        self.use_mrl = use_mrl
         self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
 
         text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
@@ -228,8 +230,14 @@ def encode_text(self, text, normalize: bool = False):
         return F.normalize(x, dim=-1) if normalize else x
 
     def forward(self, image, text):
-        image_features = self.encode_image(image, normalize=True)
-        text_features = self.encode_text(text, normalize=True)
+        if self.use_mrl:
+            # print("inside CLIP main class, using MRL")
+            image_features = self.encode_image(image, normalize=False)
+            text_features = self.encode_text(text, normalize=False)
+        else: 
+            image_features = self.encode_image(image, normalize=True)
+            text_features = self.encode_text(text, normalize=True)
+
         if self.output_dict:
             return {
                 "image_features": image_features,
diff --git a/src/training/main.py b/src/training/main.py
index f70c9f9..b0839fd 100644
--- a/src/training/main.py
+++ b/src/training/main.py
@@ -230,6 +230,7 @@ def main(args):
         image_std=args.image_std,
         aug_cfg=args.aug_cfg,
         output_dict=True,
+        use_mrl=args.force_mrl_loss
     )
     if args.distill:
         # FIXME: currenlty assumes the model your distilling from has the same tokenizer & transforms.
diff --git a/src/training/params.py b/src/training/params.py
index 36c693b..76a4c91 100644
--- a/src/training/params.py
+++ b/src/training/params.py
@@ -22,6 +22,8 @@ def __call__(self, parser, namespace, values, option_string=None):
                 kw[key] = str(value)  # fallback to string (avoid need to escape on command line)
         setattr(namespace, self.dest, kw)
 
+def parse_mrl_loss_weights(weights_str):
+    return [float(x) for x in weights_str.split(",")]
 
 def parse_args(args):
     parser = argparse.ArgumentParser()
@@ -424,6 +426,18 @@ def parse_args(args):
         default=None,
         help='Which pre-trained weights to distill from, if any.'
     )
+    parser.add_argument(
+        "--force_mrl_loss", 
+        default=False, 
+        action="store_true", 
+        help="whether to use MRL based loss"
+    )
+    parser.add_argument(
+        "--mrl_loss_weights", 
+        default=[1,1,1,1,1,1,1,1], 
+        type=parse_mrl_loss_weights, 
+        help="weights for loss weights, dimensions are considered in following order 8, 16, 32, 64, 128, 256, 512, 768"
+    )
     args = parser.parse_args(args)
 
     # If some params are not passed, we use the default values based on model name.

From a17aba1899f2d8034751166e86623ef9e9f6c02e Mon Sep 17 00:00:00 2001
From: krmayankb <krmayankb@gmail.com>
Date: Sat, 22 Jul 2023 09:11:35 -0700
Subject: [PATCH 02/18] dimensions as argument

---
 scripts/mrl_clip.sh      |  3 ++-
 src/open_clip/factory.py |  3 ++-
 src/open_clip/loss.py    | 10 ++++++----
 src/training/params.py   | 11 ++++++++++-
 4 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/scripts/mrl_clip.sh b/scripts/mrl_clip.sh
index fe6d59b..6ef7c86 100644
--- a/scripts/mrl_clip.sh
+++ b/scripts/mrl_clip.sh
@@ -8,7 +8,8 @@ torchrun --nproc_per_node 2 --master_port 3234 -m training.main \
     --gather-with-grad \
     --local-loss \
     --force_mrl_loss \
-    --mrl_loss_weights "1,1,1,1,1,1,1,1" \
+    --mrl_loss_weights "1,1,1,1,1" \
+    --mrl_dim_to_consider "768,384,192,96,48" \
     --batch-size 512 \
     --accum-freq 1 \
     --workers 4 \
diff --git a/src/open_clip/factory.py b/src/open_clip/factory.py
index f1fda90..2a27352 100644
--- a/src/open_clip/factory.py
+++ b/src/open_clip/factory.py
@@ -270,7 +270,8 @@ def create_loss(args):
             rank=args.rank,
             world_size=args.world_size,
             use_horovod=args.horovod,
-            mrl_loss_weights=args.mrl_loss_weights
+            mrl_loss_weights=args.mrl_loss_weights, 
+            dim_to_consider=args.mrl_dim_to_consider
         )
     return ClipLoss(
         local_loss=args.local_loss,
diff --git a/src/open_clip/loss.py b/src/open_clip/loss.py
index b1e9690..5a207fa 100644
--- a/src/open_clip/loss.py
+++ b/src/open_clip/loss.py
@@ -219,7 +219,8 @@ def __init__(self,
                  rank=0, 
                  world_size=1, 
                  use_horovod=False, 
-                 mrl_loss_weights = None):
+                 mrl_loss_weights = None, 
+                 dim_to_consider=None):
         super().__init__(local_loss, 
                          gather_with_grad, 
                          cache_labels, 
@@ -227,15 +228,16 @@ def __init__(self,
                          world_size, 
                          use_horovod)
         self.mrl_loss_weights = mrl_loss_weights
+        self.dim_to_consider = dim_to_consider
     
     def forward(self, image_features, text_features, logit_scale, output_dict=False):
         # print("Inside forward of MRL CLIP loss",len(self.mrl_loss_weights), self.mrl_loss_weights )
 
-        assert len(self.mrl_loss_weights) == 8, "list containing loss weights must contain 8 elements"
+        assert len(self.mrl_loss_weights) == len(self.dim_to_consider), "number of elements in loss weights and dim_to_consider should be same"
         
-        dim_to_consider = [8, 16, 32, 64, 128, 256, 512, 768]
+        # dim_to_consider = [768, 384, 192, 96, 48]
         total_loss = 0 
-        for idx, dim in enumerate(dim_to_consider): 
+        for idx, dim in enumerate(self.dim_to_consider): 
             img = F.normalize( image_features[:,:dim], dim=-1) # slice and normalize 
             text = F.normalize( text_features[:,:dim], dim=-1) # slice and normalize 
             loss = super().forward(image_features=img, text_features=text, logit_scale=logit_scale)
diff --git a/src/training/params.py b/src/training/params.py
index 76a4c91..c4ba5fd 100644
--- a/src/training/params.py
+++ b/src/training/params.py
@@ -25,6 +25,9 @@ def __call__(self, parser, namespace, values, option_string=None):
 def parse_mrl_loss_weights(weights_str):
     return [float(x) for x in weights_str.split(",")]
 
+def parse_dim_to_consider(dim_to_consider):
+    return [int(x) for x in dim_to_consider.split(",")]
+
 def parse_args(args):
     parser = argparse.ArgumentParser()
     parser.add_argument(
@@ -434,10 +437,16 @@ def parse_args(args):
     )
     parser.add_argument(
         "--mrl_loss_weights", 
-        default=[1,1,1,1,1,1,1,1], 
+        default=[1,1,1,1,1], 
         type=parse_mrl_loss_weights, 
         help="weights for loss weights, dimensions are considered in following order 8, 16, 32, 64, 128, 256, 512, 768"
     )
+    parser.add_argument(
+        "--mrl_dim_to_consider", 
+        default=[768, 384, 192, 96, 48], 
+        type=parse_dim_to_consider, 
+        help="weights for loss weights, dimensions are considered in following order 8, 16, 32, 64, 128, 256, 512, 768"
+    )
     args = parser.parse_args(args)
 
     # If some params are not passed, we use the default values based on model name.

From 9cc78dc0b6d0bb133912d505d080f95b22fd9ad1 Mon Sep 17 00:00:00 2001
From: krmayankb <krmayankb@gmail.com>
Date: Sat, 22 Jul 2023 11:09:05 -0700
Subject: [PATCH 03/18] per dim zeroshot

---
 scripts/launch_this.slurm |  4 +--
 scripts/mrl_clip.sh       | 12 ++++----
 src/training/zero_shot.py | 59 ++++++++++++++++++++++++++++-----------
 3 files changed, 50 insertions(+), 25 deletions(-)

diff --git a/scripts/launch_this.slurm b/scripts/launch_this.slurm
index 16e8be1..7e12dff 100644
--- a/scripts/launch_this.slurm
+++ b/scripts/launch_this.slurm
@@ -2,7 +2,7 @@
 
 #SBATCH --account=krishna
 #SBATCH --partition=gpu-a100
-#SBATCH --job-name=mrl
+#SBATCH --job-name=clip
 #SBATCH --error=./jobs/mrl_clip/job.err.%j  
 #SBATCH --output=./jobs/mrl_clip/job.run.%j
 #SBATCH --time=10-00:00:00
@@ -13,4 +13,4 @@
 #SBATCH --mail-user=krmayank@uw.edu
 #SBATCH --mail-type=ALL
 
-bash mrl_clip.sh
\ No newline at end of file
+bash clip.sh
\ No newline at end of file
diff --git a/scripts/mrl_clip.sh b/scripts/mrl_clip.sh
index 6ef7c86..3c0c0bb 100644
--- a/scripts/mrl_clip.sh
+++ b/scripts/mrl_clip.sh
@@ -1,5 +1,5 @@
 cd /mmfs1/gscratch/krishna/mayank/clip_clone/open_clip/src
-torchrun --nproc_per_node 2 --master_port 3234 -m training.main \
+torchrun --nproc_per_node 2 --master_port 3233 -m training.main \
     --model "ViT-B-16" \
     --train-data "/mmfs1/data/yfcc-tmp/cc_3m/train_shards/shard_{000000..003318}.tar" \
     --imagenet-val "/mmfs1/data/yfcc-tmp/imagenet/val/" \
@@ -14,10 +14,10 @@ torchrun --nproc_per_node 2 --master_port 3234 -m training.main \
     --accum-freq 1 \
     --workers 4 \
     --epochs 40 \
-    --warmup 2000 \
+    --warmup 4000 \
     --zeroshot-frequency 2 \
     --seed 0 \
-    # --report-to 'wandb' \
-    # --wandb-project-name "mrl_clip_training" \
-    # --logs "/mmfs1/gscratch/krishna/mayank/clip_clone/open_clip/src/logs/mrl_clip" \
-    # --name "mrl_clip_b512_accum_1_ep40" 
\ No newline at end of file
+    --report-to 'wandb' \
+    --wandb-project-name "mrl_clip_training" \
+    --logs "/mmfs1/gscratch/krishna/mayank/clip_clone/open_clip/src/logs/mrl_clip" \
+    --name "mrl_clip_b512_accum_1_ep40_allZeroshot" 
\ No newline at end of file
diff --git a/src/training/zero_shot.py b/src/training/zero_shot.py
index e5768b4..4243c24 100644
--- a/src/training/zero_shot.py
+++ b/src/training/zero_shot.py
@@ -9,7 +9,7 @@
 from .imagenet_zeroshot_data import imagenet_classnames, openai_imagenet_template
 
 
-def zero_shot_classifier(model, classnames, templates, args):
+def zero_shot_classifier(model, classnames, templates, dim, args):
     tokenizer = get_tokenizer(args.model)
     with torch.no_grad():
         zeroshot_weights = []
@@ -20,6 +20,7 @@ def zero_shot_classifier(model, classnames, templates, args):
                 class_embeddings = model.module.encode_text(texts)
             else:
                 class_embeddings = model.encode_text(texts)
+            class_embeddings = class_embeddings[:,:dim] # consider only dim for creating zeroshot classifier 
             class_embedding = F.normalize(class_embeddings, dim=-1).mean(dim=0)
             class_embedding /= class_embedding.norm()
             zeroshot_weights.append(class_embedding)
@@ -33,7 +34,7 @@ def accuracy(output, target, topk=(1,)):
     return [float(correct[:k].reshape(-1).float().sum(0, keepdim=True).cpu().numpy()) for k in topk]
 
 
-def run(model, classifier, dataloader, args):
+def run(model, classifier, dataloader, dim, args):
     autocast = get_autocast(args.precision)
     cast_dtype = get_cast_dtype(args.precision)
     with torch.no_grad():
@@ -50,6 +51,7 @@ def run(model, classifier, dataloader, args):
                     image_features = model.module.encode_image(images)
                 else:
                     image_features = model.encode_image(images)
+                image_features = image_features[:,:dim]
                 image_features = F.normalize(image_features, dim=-1)
                 logits = 100. * image_features @ classifier
 
@@ -74,20 +76,43 @@ def zero_shot_eval(model, data, epoch, args):
 
     logging.info('Starting zero-shot imagenet.')
 
-    logging.info('Building zero-shot classifier')
-    classifier = zero_shot_classifier(model, imagenet_classnames, openai_imagenet_template, args)
-
-    logging.info('Using classifier')
-    results = {}
-    if 'imagenet-val' in data:
-        top1, top5 = run(model, classifier, data['imagenet-val'].dataloader, args)
-        results['imagenet-zeroshot-val-top1'] = top1
-        results['imagenet-zeroshot-val-top5'] = top5
-    if 'imagenet-v2' in data:
-        top1, top5 = run(model, classifier, data['imagenet-v2'].dataloader, args)
-        results['imagenetv2-zeroshot-val-top1'] = top1
-        results['imagenetv2-zeroshot-val-top5'] = top5
-
+    # if MRL 
+    if args.force_mrl_loss:
+        results = {}
+        for dim in args.mrl_dim_to_consider: 
+            logging.info(f'Building zero-shot classifier dim-{dim}')
+            classifier = zero_shot_classifier(model, imagenet_classnames, openai_imagenet_template, dim, args)
+
+            logging.info('Using classifier')
+            if 'imagenet-val' in data:
+                top1, top5 = run(model, classifier, data['imagenet-val'].dataloader, dim, args)
+                top1_name = f'imagenet-zeroshot-val-d{dim}-top1'
+                top5_name = f'imagenet-zeroshot-val-d{dim}-top5'
+                results[top1_name] = top1
+                results[top5_name] = top5
+            if 'imagenet-v2' in data:
+                top1, top5 = run(model, classifier, data['imagenet-v2'].dataloader, dim, args)
+                top1_name = f'imagenet-zeroshot-val-d{dim}-top1'
+                top5_name = f'imagenet-zeroshot-val-d{dim}-top5'
+                results[top1_name] = top1
+                results[top5_name] = top5
+    
+    # for other losses than MRL
+    else: 
+        logging.info('Building zero-shot classifier')
+        classifier = zero_shot_classifier(model, imagenet_classnames, openai_imagenet_template, args)
+
+        logging.info('Using classifier')
+        results = {}
+        if 'imagenet-val' in data:
+            top1, top5 = run(model, classifier, data['imagenet-val'].dataloader, args)
+            results['imagenet-zeroshot-val-top1'] = top1
+            results['imagenet-zeroshot-val-top5'] = top5
+        if 'imagenet-v2' in data:
+            top1, top5 = run(model, classifier, data['imagenet-v2'].dataloader, args)
+            results['imagenetv2-zeroshot-val-top1'] = top1
+            results['imagenetv2-zeroshot-val-top5'] = top5
+
+    # if not MRL 
     logging.info('Finished zero-shot imagenet.')
-
     return results

From a19ad920d12213ec24c34bfb82ef8900ec6cc2d0 Mon Sep 17 00:00:00 2001
From: krmayankb <krmayankb@gmail.com>
Date: Mon, 24 Jul 2023 11:36:54 -0700
Subject: [PATCH 04/18] generalized zeroshot

---
 scripts/clip.sh           | 20 ++++++++++++++++++++
 src/training/zero_shot.py | 12 +++++++-----
 2 files changed, 27 insertions(+), 5 deletions(-)
 create mode 100644 scripts/clip.sh

diff --git a/scripts/clip.sh b/scripts/clip.sh
new file mode 100644
index 0000000..b878d08
--- /dev/null
+++ b/scripts/clip.sh
@@ -0,0 +1,20 @@
+cd /mmfs1/gscratch/krishna/mayank/clip_clone/open_clip/src
+torchrun --nproc_per_node 2 --master_port 3467 -m training.main \
+    --model "ViT-B-16" \
+    --train-data "/mmfs1/data/yfcc-tmp/cc_3m/train_shards/shard_{000000..003318}.tar" \
+    --imagenet-val "/mmfs1/data/yfcc-tmp/imagenet/val/" \
+    --dataset-type webdataset \
+    --precision amp \
+    --gather-with-grad \
+    --local-loss \
+    --batch-size 512 \
+    --accum-freq 1 \
+    --workers 4 \
+    --epochs 40 \
+    --warmup 4000 \
+    --zeroshot-frequency 2 \
+    --seed 0 \
+    --report-to 'wandb' \
+    --wandb-project-name "mrl_clip_training" \
+    --logs "/mmfs1/gscratch/krishna/mayank/clip_clone/open_clip/src/logs/mrl_clip" \
+    --name "clip_b512_accum_1_ep40_bugfixed" 
\ No newline at end of file
diff --git a/src/training/zero_shot.py b/src/training/zero_shot.py
index 4243c24..ac814c8 100644
--- a/src/training/zero_shot.py
+++ b/src/training/zero_shot.py
@@ -20,7 +20,8 @@ def zero_shot_classifier(model, classnames, templates, dim, args):
                 class_embeddings = model.module.encode_text(texts)
             else:
                 class_embeddings = model.encode_text(texts)
-            class_embeddings = class_embeddings[:,:dim] # consider only dim for creating zeroshot classifier 
+            if args.force_mrl_loss:
+                class_embeddings = class_embeddings[:,:dim] # consider only dim for creating zeroshot classifier 
             class_embedding = F.normalize(class_embeddings, dim=-1).mean(dim=0)
             class_embedding /= class_embedding.norm()
             zeroshot_weights.append(class_embedding)
@@ -51,7 +52,8 @@ def run(model, classifier, dataloader, dim, args):
                     image_features = model.module.encode_image(images)
                 else:
                     image_features = model.encode_image(images)
-                image_features = image_features[:,:dim]
+                if args.force_mrl_loss:
+                    image_features = image_features[:,:dim]
                 image_features = F.normalize(image_features, dim=-1)
                 logits = 100. * image_features @ classifier
 
@@ -100,16 +102,16 @@ def zero_shot_eval(model, data, epoch, args):
     # for other losses than MRL
     else: 
         logging.info('Building zero-shot classifier')
-        classifier = zero_shot_classifier(model, imagenet_classnames, openai_imagenet_template, args)
+        classifier = zero_shot_classifier(model, imagenet_classnames, openai_imagenet_template, 0, args)
 
         logging.info('Using classifier')
         results = {}
         if 'imagenet-val' in data:
-            top1, top5 = run(model, classifier, data['imagenet-val'].dataloader, args)
+            top1, top5 = run(model, classifier, data['imagenet-val'].dataloader, 0, args)
             results['imagenet-zeroshot-val-top1'] = top1
             results['imagenet-zeroshot-val-top5'] = top5
         if 'imagenet-v2' in data:
-            top1, top5 = run(model, classifier, data['imagenet-v2'].dataloader, args)
+            top1, top5 = run(model, classifier, data['imagenet-v2'].dataloader, 0, args)
             results['imagenetv2-zeroshot-val-top1'] = top1
             results['imagenetv2-zeroshot-val-top5'] = top5
 

From fc145863a17174cdfec80a0ea57598995b48465a Mon Sep 17 00:00:00 2001
From: krmayankb <krmayankb@gmail.com>
Date: Tue, 25 Jul 2023 21:33:56 -0700
Subject: [PATCH 05/18] logit scale per granularity

comment
---
 scripts/finetuning_original.sh | 25 ++++++++++++++++++++
 scripts/finetuning_weighted.sh | 25 ++++++++++++++++++++
 scripts/launch_this.slurm      |  4 ++--
 scripts/mrl_clip.sh            |  4 ++--
 scripts/resume.sh              | 25 ++++++++++++++++++++
 src/open_clip/factory.py       | 16 +++++++++----
 src/open_clip/loss.py          | 14 +++++++-----
 src/open_clip/model.py         | 34 +++++++++++++++++++--------
 src/training/main.py           |  3 ++-
 src/training/train.py          | 42 +++++++++++++++++++++++++---------
 10 files changed, 155 insertions(+), 37 deletions(-)
 create mode 100644 scripts/finetuning_original.sh
 create mode 100644 scripts/finetuning_weighted.sh
 create mode 100644 scripts/resume.sh

diff --git a/scripts/finetuning_original.sh b/scripts/finetuning_original.sh
new file mode 100644
index 0000000..9382e4b
--- /dev/null
+++ b/scripts/finetuning_original.sh
@@ -0,0 +1,25 @@
+cd /mmfs1/gscratch/krishna/mayank/clip_clone/open_clip/src
+torchrun --nproc_per_node 2 --master_port 4556 -m training.main \
+    --model "ViT-B-16" \
+    --pretrained "laion400m_e32" \
+    --train-data "/mmfs1/data/yfcc-tmp/cc_3m/train_shards/shard_{000000..003318}.tar" \
+    --imagenet-val "/mmfs1/data/yfcc-tmp/imagenet/val/" \
+    --dataset-type webdataset \
+    --precision amp \
+    --gather-with-grad \
+    --local-loss \
+    --force_mrl_loss \
+    --mrl_loss_weights "1,1,1,1,1" \
+    --mrl_dim_to_consider "768,384,192,96,48" \
+    --accum-freq 1 \
+    --batch-size 512 \
+    --lr 1e-07 \
+    --workers 4 \
+    --epochs 10 \
+    --warmup 500 \
+    --zeroshot-frequency 1 \
+    --seed 0 \
+    --report-to 'wandb' \
+    --wandb-project-name "mrl_clip_training" \
+    --logs "/mmfs1/gscratch/krishna/mayank/clip_clone/open_clip/src/logs/mrl_clip" \
+    --name "ViT-B-16_liaon400m_e32_finetune_mrl_ep10_warmup_500_lr1e-07" 
\ No newline at end of file
diff --git a/scripts/finetuning_weighted.sh b/scripts/finetuning_weighted.sh
new file mode 100644
index 0000000..6b5ada3
--- /dev/null
+++ b/scripts/finetuning_weighted.sh
@@ -0,0 +1,25 @@
+cd /mmfs1/gscratch/krishna/mayank/clip_clone/open_clip/src
+torchrun --nproc_per_node 2 --master_port 4534 -m training.main \
+    --model "ViT-B-32" \
+    --pretrained "laion400m_e32" \
+    --train-data "/mmfs1/data/yfcc-tmp/cc_3m/train_shards/shard_{000000..003318}.tar" \
+    --imagenet-val "/mmfs1/data/yfcc-tmp/imagenet/val/" \
+    --dataset-type webdataset \
+    --precision amp \
+    --gather-with-grad \
+    --local-loss \
+    --force_mrl_loss \
+    --mrl_loss_weights "0.3,0.25,0.2,0.15,0.1" \
+    --mrl_dim_to_consider "768,384,192,96,48" \
+    --accum-freq 1 \
+    --batch-size 512 \
+    --lr 1e-07 \
+    --workers 4 \
+    --epochs 10 \
+    --warmup 500 \
+    --zeroshot-frequency 1 \
+    --seed 0 \
+    --report-to 'wandb' \
+    --wandb-project-name "mrl_clip_training" \
+    --logs "/mmfs1/gscratch/krishna/mayank/clip_clone/open_clip/src/logs/mrl_clip" \
+    --name "ViT-B-16_liaon400m_e32_finetune_mrl_ep10_warmup_500_wl030250201501_lr1e-07" 
\ No newline at end of file
diff --git a/scripts/launch_this.slurm b/scripts/launch_this.slurm
index 7e12dff..41dce10 100644
--- a/scripts/launch_this.slurm
+++ b/scripts/launch_this.slurm
@@ -2,7 +2,7 @@
 
 #SBATCH --account=krishna
 #SBATCH --partition=gpu-a100
-#SBATCH --job-name=clip
+#SBATCH --job-name=MClip1
 #SBATCH --error=./jobs/mrl_clip/job.err.%j  
 #SBATCH --output=./jobs/mrl_clip/job.run.%j
 #SBATCH --time=10-00:00:00
@@ -13,4 +13,4 @@
 #SBATCH --mail-user=krmayank@uw.edu
 #SBATCH --mail-type=ALL
 
-bash clip.sh
\ No newline at end of file
+bash mrl_clip.sh
\ No newline at end of file
diff --git a/scripts/mrl_clip.sh b/scripts/mrl_clip.sh
index 3c0c0bb..f991e61 100644
--- a/scripts/mrl_clip.sh
+++ b/scripts/mrl_clip.sh
@@ -8,7 +8,7 @@ torchrun --nproc_per_node 2 --master_port 3233 -m training.main \
     --gather-with-grad \
     --local-loss \
     --force_mrl_loss \
-    --mrl_loss_weights "1,1,1,1,1" \
+    --mrl_loss_weights "0.1,0.15,0.2,0.25,0.3" \
     --mrl_dim_to_consider "768,384,192,96,48" \
     --batch-size 512 \
     --accum-freq 1 \
@@ -20,4 +20,4 @@ torchrun --nproc_per_node 2 --master_port 3233 -m training.main \
     --report-to 'wandb' \
     --wandb-project-name "mrl_clip_training" \
     --logs "/mmfs1/gscratch/krishna/mayank/clip_clone/open_clip/src/logs/mrl_clip" \
-    --name "mrl_clip_b512_accum_1_ep40_allZeroshot" 
\ No newline at end of file
+    --name "mrl_clip_b512_accum_1_ep40_diffLogitScale_D082723_wl010150202503" 
\ No newline at end of file
diff --git a/scripts/resume.sh b/scripts/resume.sh
new file mode 100644
index 0000000..194c212
--- /dev/null
+++ b/scripts/resume.sh
@@ -0,0 +1,25 @@
+cd /mmfs1/gscratch/krishna/mayank/clip_clone/open_clip/src
+torchrun --nproc_per_node 2 -m --master_port 3436 training.main \
+    --model "ViT-B-16" \
+    --train-data "/mmfs1/data/yfcc-tmp/cc_3m/train_shards/shard_{000000..003318}.tar" \
+    --imagenet-val "/mmfs1/data/yfcc-tmp/imagenet/val/" \
+    --dataset-type webdataset \
+    --precision amp \
+    --gather-with-grad \
+    --local-loss \
+    --accum-freq 1 \
+    --workers 4 \
+    --epochs 60 \
+    --warmup 1000 \
+    --zeroshot-frequency 1 \
+    --seed 0 \
+    --resume "/mmfs1/gscratch/krishna/mayank/clip_clone/open_clip/src/logs/mrl_clip/clip_b512_accum_1_ep40_bugfixed/checkpoints/epoch_40.pt" \
+    --report-to 'wandb' \
+    --wandb-project-name "mrl_clip_training" \
+    --logs "/mmfs1/gscratch/krishna/mayank/clip_clone/open_clip/src/logs/mrl_clip" \
+    --name "clip_b512_accum_1_ep40_resumefrom40" 
+
+
+    # --force_mrl_loss \
+    # --mrl_loss_weights "1,1,1,1,1" \
+    # --mrl_dim_to_consider "768,384,192,96,48" \
\ No newline at end of file
diff --git a/src/open_clip/factory.py b/src/open_clip/factory.py
index 2a27352..b23d2bf 100644
--- a/src/open_clip/factory.py
+++ b/src/open_clip/factory.py
@@ -119,7 +119,8 @@ def create_model(
         cache_dir: Optional[str] = None,
         output_dict: Optional[bool] = None,
         require_pretrained: bool = False,
-        use_mrl: bool = False
+        use_mrl: bool = False, 
+        mrl_dim: int = 0
 ):
     has_hf_hub_prefix = model_name.startswith(HF_HUB_PREFIX)
     if has_hf_hub_prefix:
@@ -192,7 +193,7 @@ def create_model(
             else:
                 model = CustomTextCLIP(**model_cfg, cast_dtype=cast_dtype)
         else:
-            model = CLIP(**model_cfg, cast_dtype=cast_dtype, use_mrl=use_mrl) # intialise with mrl based config if use_mrl
+            model = CLIP(**model_cfg, cast_dtype=cast_dtype, use_mrl=use_mrl, mrl_dim=mrl_dim) # intialise with mrl based config if use_mrl
 
         pretrained_loaded = False
         if pretrained:
@@ -205,7 +206,10 @@ def create_model(
 
             if checkpoint_path:
                 logging.info(f'Loading pretrained {model_name} weights ({pretrained}).')
-                load_checkpoint(model, checkpoint_path)
+                if use_mrl: 
+                    load_checkpoint(model, checkpoint_path, strict=False)
+                else: 
+                    load_checkpoint(model, checkpoint_path)
             else:
                 error_str = (
                     f'Pretrained weights ({pretrained}) not found for model {model_name}.'
@@ -300,7 +304,8 @@ def create_model_and_transforms(
         aug_cfg: Optional[Union[Dict[str, Any], AugmentationCfg]] = None,
         cache_dir: Optional[str] = None,
         output_dict: Optional[bool] = None,
-        use_mrl = False
+        use_mrl: Optional[bool] = False, 
+        mrl_dim: int = 0
 ):
     model = create_model(
         model_name,
@@ -316,7 +321,8 @@ def create_model_and_transforms(
         pretrained_hf=pretrained_hf,
         cache_dir=cache_dir,
         output_dict=output_dict,
-        use_mrl=use_mrl
+        use_mrl=use_mrl, 
+        mrl_dim = mrl_dim
     )
 
     image_mean = image_mean or getattr(model.visual, 'image_mean', None)
diff --git a/src/open_clip/loss.py b/src/open_clip/loss.py
index 5a207fa..831caaa 100644
--- a/src/open_clip/loss.py
+++ b/src/open_clip/loss.py
@@ -236,14 +236,16 @@ def forward(self, image_features, text_features, logit_scale, output_dict=False)
         assert len(self.mrl_loss_weights) == len(self.dim_to_consider), "number of elements in loss weights and dim_to_consider should be same"
         
         # dim_to_consider = [768, 384, 192, 96, 48]
-        total_loss = 0 
+        # total_loss = 0
+        loss_list = []
         for idx, dim in enumerate(self.dim_to_consider): 
             img = F.normalize( image_features[:,:dim], dim=-1) # slice and normalize 
             text = F.normalize( text_features[:,:dim], dim=-1) # slice and normalize 
-            loss = super().forward(image_features=img, text_features=text, logit_scale=logit_scale)
-            total_loss += self.mrl_loss_weights[idx] * loss
-        
+            loss = super().forward(image_features=img, text_features=text, logit_scale=logit_scale[idx])
+            # total_loss += self.mrl_loss_weights[idx] * loss
+            loss_list.append(self.mrl_loss_weights[idx] * loss)
+            
         if output_dict:
-            return {"mrl_clip_loss": total_loss}        
+            return {f"mrl_clip_loss_{key}": value for key, value in zip(self.dim_to_consider, loss_list)} 
         
-        return total_loss
+        return loss_list
diff --git a/src/open_clip/model.py b/src/open_clip/model.py
index 7729e14..942e102 100644
--- a/src/open_clip/model.py
+++ b/src/open_clip/model.py
@@ -184,11 +184,13 @@ def __init__(
             quick_gelu: bool = False,
             cast_dtype: Optional[torch.dtype] = None,
             output_dict: bool = False,
-            use_mrl: bool = False
+            use_mrl: bool = False, 
+            mrl_dim: int = 0,
     ):
         super().__init__()
         self.output_dict = output_dict
         self.use_mrl = use_mrl
+        self.mrl_dim = mrl_dim
         self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
 
         text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
@@ -199,8 +201,12 @@ def __init__(
         self.ln_final = text.ln_final
         self.text_projection = text.text_projection
         self.register_buffer('attn_mask', text.attn_mask, persistent=False)
-
-        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+        
+        if use_mrl:
+            # print("using mrl inside clip initialization")
+            self.logit_scale = nn.ParameterList([nn.Parameter(torch.ones([]) * np.log(1 / 0.07)) for _ in range(self.mrl_dim)])
+        else: 
+            self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
 
     def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
         # lock image tower as per LiT - https://arxiv.org/abs/2111.07991
@@ -234,17 +240,25 @@ def forward(self, image, text):
             # print("inside CLIP main class, using MRL")
             image_features = self.encode_image(image, normalize=False)
             text_features = self.encode_text(text, normalize=False)
+
+            if self.output_dict: 
+                return {
+                    "image_features": image_features,
+                    "text_features": text_features,
+                    "logit_scale": [scale.exp() for scale in self.logit_scale]
+                }
+            return image_features, text_features, [scale.exp() for scale in self.logit_scale]
         else: 
             image_features = self.encode_image(image, normalize=True)
             text_features = self.encode_text(text, normalize=True)
+            if self.output_dict:
+                return {
+                    "image_features": image_features,
+                    "text_features": text_features,
+                    "logit_scale": self.logit_scale.exp()
+                }
+            return image_features, text_features, self.logit_scale.exp()
 
-        if self.output_dict:
-            return {
-                "image_features": image_features,
-                "text_features": text_features,
-                "logit_scale": self.logit_scale.exp()
-            }
-        return image_features, text_features, self.logit_scale.exp()
 
 
 class CustomTextCLIP(nn.Module):
diff --git a/src/training/main.py b/src/training/main.py
index b0839fd..647ca02 100644
--- a/src/training/main.py
+++ b/src/training/main.py
@@ -230,7 +230,8 @@ def main(args):
         image_std=args.image_std,
         aug_cfg=args.aug_cfg,
         output_dict=True,
-        use_mrl=args.force_mrl_loss
+        use_mrl=args.force_mrl_loss,
+        mrl_dim = len(args.mrl_dim_to_consider)
     )
     if args.distill:
         # FIXME: currenlty assumes the model your distilling from has the same tokenizer & transforms.
diff --git a/src/training/train.py b/src/training/train.py
index e0a140f..86b67db 100644
--- a/src/training/train.py
+++ b/src/training/train.py
@@ -174,7 +174,11 @@ def train_one_epoch(model, data, loss, epoch, optimizer, scaler, scheduler, dist
 
         # Note: we clamp to 4.6052 = ln(100), as in the original paper.
         with torch.no_grad():
-            unwrap_model(model).logit_scale.clamp_(0, math.log(100))
+            if args.force_mrl_loss: 
+                for idx in range(len(args.mrl_dim_to_consider)):
+                    unwrap_model(model).logit_scale[idx].clamp_(0, math.log(100))
+            else: 
+                unwrap_model(model).logit_scale.clamp_(0, math.log(100))
 
         batch_time_m.update(time.time() - end)
         end = time.time()
@@ -191,7 +195,12 @@ def train_one_epoch(model, data, loss, epoch, optimizer, scaler, scheduler, dist
                     losses_m[key] = AverageMeter()
                 losses_m[key].update(val.item(), batch_size)
 
-            logit_scale_scalar = logit_scale.item()
+            if args.force_mrl_loss: 
+                logit_scale_scalar = [logit_scale[i].item() for i in range(len(args.mrl_dim_to_consider))]
+                logit_scale_string = " ,".join([f"{item:.3f}" for item in logit_scale_scalar])
+            else: 
+                logit_scale_scalar = logit_scale.item()
+                logit_scale_string = f"{logit_scale_scalar:.3f}"
             loss_log = " ".join(
                 [
                     f"{loss_name.capitalize()}: {loss_m.val:#.5g} ({loss_m.avg:#.5g})" 
@@ -205,18 +214,29 @@ def train_one_epoch(model, data, loss, epoch, optimizer, scaler, scheduler, dist
                 f"Data (t): {data_time_m.avg:.3f} "
                 f"Batch (t): {batch_time_m.avg:.3f}, {samples_per_second:#g}/s, {samples_per_second_per_gpu:#g}/s/gpu "
                 f"LR: {optimizer.param_groups[0]['lr']:5f} "
-                f"Logit Scale: {logit_scale_scalar:.3f} " + loss_log
+                f"Logit Scale: {logit_scale_string} " + loss_log
             )
 
             # Save train loss / etc. Using non avg meter values as loggers have their own smoothing
-            log_data = {
-                "data_time": data_time_m.val,
-                "batch_time": batch_time_m.val,
-                "samples_per_second": samples_per_second,
-                "samples_per_second_per_gpu": samples_per_second_per_gpu,
-                "scale": logit_scale_scalar,
-                "lr": optimizer.param_groups[0]["lr"]
-            }            
+            if args.force_mrl_loss: 
+                log_data = {
+                    "data_time": data_time_m.val,
+                    "batch_time": batch_time_m.val,
+                    "samples_per_second": samples_per_second,
+                    "samples_per_second_per_gpu": samples_per_second_per_gpu,
+                    "lr": optimizer.param_groups[0]["lr"]
+                }            
+                log_data.update({f"scale_{dim}": logit_scale_scalar[idx] for idx, dim in enumerate(args.mrl_dim_to_consider)})
+            else:
+                log_data = {
+                    "data_time": data_time_m.val,
+                    "batch_time": batch_time_m.val,
+                    "samples_per_second": samples_per_second,
+                    "samples_per_second_per_gpu": samples_per_second_per_gpu,
+                    "scale": logit_scale_scalar,
+                    "lr": optimizer.param_groups[0]["lr"]
+                }            
+
             log_data.update({name:val.val for name,val in losses_m.items()})
 
             for name, val in log_data.items():

From a86610f0d9e89fcb1301a28bcf62cd731e3bd2d7 Mon Sep 17 00:00:00 2001
From: krmayankb <krmayankb@gmail.com>
Date: Sat, 16 Sep 2023 21:16:14 -0700
Subject: [PATCH 06/18] tpu support

---
 src/open_clip/loss.py       | 19 ++++++++++++++++---
 src/open_clip/model.py      | 15 +++++++++++++++
 src/training/distributed.py | 11 ++++++++++-
 src/training/main.py        | 34 ++++++++++++++++++++++++++++++----
 src/training/params.py      |  6 ++++++
 src/training/train.py       | 35 +++++++++++++++++++++++++++++------
 6 files changed, 106 insertions(+), 14 deletions(-)

diff --git a/src/open_clip/loss.py b/src/open_clip/loss.py
index 831caaa..681f071 100644
--- a/src/open_clip/loss.py
+++ b/src/open_clip/loss.py
@@ -15,6 +15,11 @@
 except ImportError:
     hvd = None
 
+try:
+    import torch_xla.core.xla_model as xm
+except ImportError:
+    xm = None
+
 
 def gather_features(
         image_features,
@@ -230,6 +235,14 @@ def __init__(self,
         self.mrl_loss_weights = mrl_loss_weights
         self.dim_to_consider = dim_to_consider
     
+    def normalize(self, features, dim=-1):
+        if xm.xla_device():
+            norm = xm.all_reduce("sum", features ** 2)
+            norm = torch.sqrt(norm)
+            features = features / norm   
+            return features
+        return F.normalize(features, dim=dim)
+    
     def forward(self, image_features, text_features, logit_scale, output_dict=False):
         # print("Inside forward of MRL CLIP loss",len(self.mrl_loss_weights), self.mrl_loss_weights )
 
@@ -239,9 +252,9 @@ def forward(self, image_features, text_features, logit_scale, output_dict=False)
         # total_loss = 0
         loss_list = []
         for idx, dim in enumerate(self.dim_to_consider): 
-            img = F.normalize( image_features[:,:dim], dim=-1) # slice and normalize 
-            text = F.normalize( text_features[:,:dim], dim=-1) # slice and normalize 
-            loss = super().forward(image_features=img, text_features=text, logit_scale=logit_scale[idx])
+            img = self.normalize(image_features[:,:dim], dim=-1) # slice and normalize 
+            txt = self.normalize(text_features[:,:dim], dim=-1) # slice and normalize 
+            loss = super().forward(image_features=img, text_features=txt, logit_scale=logit_scale[idx])
             # total_loss += self.mrl_loss_weights[idx] * loss
             loss_list.append(self.mrl_loss_weights[idx] * loss)
             
diff --git a/src/open_clip/model.py b/src/open_clip/model.py
index 942e102..9f7ce51 100644
--- a/src/open_clip/model.py
+++ b/src/open_clip/model.py
@@ -19,6 +19,10 @@
 from .transformer import LayerNormFp32, LayerNorm, QuickGELU, Attention, VisionTransformer, TextTransformer
 from .utils import to_2tuple
 
+try: 
+    import torch_xla.core.xla_model as xm
+except ImportError:
+    pass
 
 @dataclass
 class CLIPVisionCfg:
@@ -217,8 +221,17 @@ def set_grad_checkpointing(self, enable=True):
         self.visual.set_grad_checkpointing(enable)
         self.transformer.grad_checkpointing = enable
 
+    def xm_normalize(self, features, normalize: bool=False):
+        if normalize:
+            norm = xm.all_reduce("sum", features ** 2)
+            norm = torch.sqrt(norm)
+            features = features / norm    
+        return features
+
     def encode_image(self, image, normalize: bool = False):
         features = self.visual(image)
+        if xm.xla_device():
+            return self.xm_normalize(features, normalize=normalize)
         return F.normalize(features, dim=-1) if normalize else features
 
     def encode_text(self, text, normalize: bool = False):
@@ -233,6 +246,8 @@ def encode_text(self, text, normalize: bool = False):
         x = self.ln_final(x)  # [batch_size, n_ctx, transformer.width]
         # take features from the eot embedding (eot_token is the highest number in each sequence)
         x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+        if xm.xla_device():
+            return self.xm_normalize(x, normalize=normalize)
         return F.normalize(x, dim=-1) if normalize else x
 
     def forward(self, image, text):
diff --git a/src/training/distributed.py b/src/training/distributed.py
index 268a6c7..c120f1b 100644
--- a/src/training/distributed.py
+++ b/src/training/distributed.py
@@ -8,6 +8,10 @@
 except ImportError:
     hvd = None
 
+try: 
+    import torch_xla.core.xla_model as xm
+except ImportError:
+    xm = None
 
 def is_global_master(args):
     return args.rank == 0
@@ -39,7 +43,6 @@ def is_using_distributed():
         return int(os.environ['SLURM_NTASKS']) > 1
     return False
 
-
 def world_info_from_env():
     local_rank = 0
     for v in ('LOCAL_RANK', 'MPI_LOCALRANKID', 'SLURM_LOCALID', 'OMPI_COMM_WORLD_LOCAL_RANK'):
@@ -107,6 +110,8 @@ def init_distributed_device(args):
         else:
             device = 'cuda:0'
         torch.cuda.set_device(device)
+    elif xm is not None:
+        device = xm.xla_device()
     else:
         device = 'cpu'
     args.device = device
@@ -118,6 +123,8 @@ def broadcast_object(args, obj, src=0):
     # broadcast a pickle-able python object from rank-0 to all ranks
     if args.horovod:
         return hvd.broadcast_object(obj, root_rank=src)
+    elif xm is not None:
+        return xm.mesh_reduce('broadcast_object', obj, lambda x: x, src)
     else:
         if args.rank == src:
             objects = [obj]
@@ -131,6 +138,8 @@ def all_gather_object(args, obj, dst=0):
     # gather a pickle-able python object across all ranks
     if args.horovod:
         return hvd.allgather_object(obj)
+    elif xm is not None:
+        return xm.mesh_reduce('all_gather_object', obj, lambda x: x, dst)
     else:
         objects = [None for _ in range(args.world_size)]
         dist.all_gather_object(objects, obj)
diff --git a/src/training/main.py b/src/training/main.py
index 647ca02..388fe1c 100644
--- a/src/training/main.py
+++ b/src/training/main.py
@@ -1,3 +1,5 @@
+# modify this file to support distributed training on tpus and gpus
+# this file is a modified version of the original main.py file from the open_clip repo
 import glob
 import logging
 import os
@@ -12,6 +14,15 @@
 from torch import optim
 from torch.cuda.amp import GradScaler
 
+try:
+    import torch_xla.core.xla_model as xm
+    import torch_xla.distributed.parallel_loader as pl
+    import torch_xla.distributed.xla_multiprocessing as xmp
+except ImportError:
+    xm = None
+    pl = None
+    xmp = None
+
 try:
     import wandb
 except ImportError:
@@ -69,7 +80,7 @@ def get_latest_checkpoint(path: str, remote : bool):
 
 def main(args):
     args = parse_args(args)
-
+    # check if args.use_tpu is set, if so, use torch_xla        
     if torch.cuda.is_available():
         # This enables tf32 on Ampere GPUs which is only 8% slower than
         # float16 and almost as accurate as float32
@@ -77,6 +88,11 @@ def main(args):
         torch.backends.cuda.matmul.allow_tf32 = True
         torch.backends.cudnn.benchmark = True
         torch.backends.cudnn.deterministic = False
+    elif args.use_tpu:
+        assert xm is not None, "Please install torch_xla to use TPUs."
+        xm.set_rng_state(args.seed, device=xm.xla_device())
+        torch.backends.cudnn.benchmark = True
+        torch.backends.cudnn.deterministic = False
 
     # fully initialize distributed device environment
     device = init_distributed_device(args)
@@ -312,7 +328,7 @@ def main(args):
             hvd.broadcast_parameters(model.state_dict(), root_rank=0)
             hvd.broadcast_optimizer_state(optimizer, root_rank=0)
 
-        scaler = GradScaler() if args.precision == "amp" else None
+        scaler = GradScaler() if args.precision == "amp" and not xm.xla_device() else None
 
     # optionally resume from a checkpoint
     start_epoch = 0
@@ -336,7 +352,14 @@ def main(args):
             logging.info(f"=> loaded checkpoint '{args.resume}' (epoch {start_epoch})")
 
     # initialize datasets
-    data = get_data(args, (preprocess_train, preprocess_val), epoch=start_epoch, tokenizer=get_tokenizer(args.model))
+    if args.use_tpu:
+        if not xm.is_master_ordinal():
+            xm.rendezvous('download_only_once')
+        data = get_data(args, (preprocess_train, preprocess_val), epoch=start_epoch, tokenizer=get_tokenizer(args.model))      
+        if xm.is_master_ordinal():
+            xm.rendezvous('download_only_once')  
+    else:
+        data = get_data(args, (preprocess_train, preprocess_val), epoch=start_epoch, tokenizer=get_tokenizer(args.model))
     assert len(data), 'At least one train or eval dataset must be specified.'
 
     # create scheduler if train
@@ -469,4 +492,7 @@ def copy_codebase(args):
 
 
 if __name__ == "__main__":
-    main(sys.argv[1:])
+    if xmp is not None:
+        xmp.spawn(main(sys.argv[1:]), nprocs=8, start_method='fork')
+    else:
+        main(sys.argv[1:])
diff --git a/src/training/params.py b/src/training/params.py
index c4ba5fd..f480b20 100644
--- a/src/training/params.py
+++ b/src/training/params.py
@@ -447,6 +447,12 @@ def parse_args(args):
         type=parse_dim_to_consider, 
         help="weights for loss weights, dimensions are considered in following order 8, 16, 32, 64, 128, 256, 512, 768"
     )
+    parser.add_argument(
+        "--use_tpu",
+        default=False, 
+        action="store_true",
+        help="Weather to use TPUs for training"
+    )
     args = parser.parse_args(args)
 
     # If some params are not passed, we use the default values based on model name.
diff --git a/src/training/train.py b/src/training/train.py
index 86b67db..b2d5a50 100644
--- a/src/training/train.py
+++ b/src/training/train.py
@@ -8,12 +8,20 @@
 import torch
 import torch.nn.functional as F
 from torch.nn.parallel.distributed import DistributedDataParallel
+import contextlib
 
 try:
     import wandb
 except ImportError:
     wandb = None
 
+try:
+    import torch_xla.core.xla_model as xm
+    import torch_xla.distributed.parallel_loader as pl
+except ImportError:
+    xm = None
+    pl = None
+
 from open_clip import get_cast_dtype, CLIP, CustomTextCLIP
 from .distributed import is_master
 from .zero_shot import zero_shot_eval
@@ -61,7 +69,7 @@ def backward(total_loss, scaler):
 
 def train_one_epoch(model, data, loss, epoch, optimizer, scaler, scheduler, dist_model, args, tb_writer=None):
     device = torch.device(args.device)
-    autocast = get_autocast(args.precision)
+    autocast = get_autocast(args.precision) if not args.use_tpu else contextlib.nullcontext
     cast_dtype = get_cast_dtype(args.precision)
 
 
@@ -72,7 +80,10 @@ def train_one_epoch(model, data, loss, epoch, optimizer, scaler, scheduler, dist
     data['train'].set_epoch(epoch)  # set epoch in process safe manner via sampler or shared_epoch
     dataloader = data['train'].dataloader
     num_batches_per_epoch = dataloader.num_batches // args.accum_freq
+    samples_per_epoch = dataloader.num_samples
     sample_digits = math.ceil(math.log(dataloader.num_samples + 1, 10))
+    if args.use_tpu:
+        dataloader = pl.ParallelLoader(dataloader, [device]).per_device_loader(device)
 
     if args.accum_freq > 1:
         accum_images, accum_texts, accum_features = [], [], {}
@@ -166,7 +177,10 @@ def train_one_epoch(model, data, loss, epoch, optimizer, scaler, scheduler, dist
         else:
             if args.grad_clip_norm is not None:
                 torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip_norm, norm_type=2.0)
-            optimizer.step()
+            if args.use_tpu:
+                xm.optimizer_step(optimizer)
+            else:
+                optimizer.step()
 
         # reset gradient accum, if enabled
         if args.accum_freq > 1:
@@ -186,14 +200,17 @@ def train_one_epoch(model, data, loss, epoch, optimizer, scaler, scheduler, dist
         if is_master(args) and (i_accum % args.log_every_n_steps == 0 or batch_count == num_batches_per_epoch):
             batch_size = len(images)
             num_samples = batch_count * batch_size * args.accum_freq * args.world_size
-            samples_per_epoch = dataloader.num_samples
+            # samples_per_epoch = dataloader.num_samples
             percent_complete = 100.0 * batch_count / num_batches_per_epoch
 
             # NOTE loss is coarsely sampled, just master node and per log update
             for key, val in losses.items():
                 if key not in losses_m:
                     losses_m[key] = AverageMeter()
-                losses_m[key].update(val.item(), batch_size)
+                if args.use_tpu:
+                    losses_m[key].update(val, batch_size)
+                else: 
+                    losses_m[key].update(val.item(), batch_size)
 
             if args.force_mrl_loss: 
                 logit_scale_scalar = [logit_scale[i].item() for i in range(len(args.mrl_dim_to_consider))]
@@ -263,7 +280,7 @@ def evaluate(model, data, epoch, args, tb_writer=None):
     zero_shot_metrics = zero_shot_eval(model, data, epoch, args)
     metrics.update(zero_shot_metrics)
 
-    autocast = get_autocast(args.precision)
+    autocast = get_autocast(args.precision) if not args.use_tpu else contextlib.nullcontext
     cast_dtype = get_cast_dtype(args.precision)
 
     if 'val' in data and (args.val_frequency and ((epoch % args.val_frequency) == 0 or epoch == args.epochs)):
@@ -271,6 +288,9 @@ def evaluate(model, data, epoch, args, tb_writer=None):
         num_samples = 0
         samples_per_val = dataloader.num_samples
 
+        if args.use_tpu:
+            dataloader = pl.ParallelLoader(dataloader, [device]).per_device_loader(device)
+
         # FIXME this does not scale past small eval datasets
         # all_image_features @ all_text_features will blow up memory and compute very quickly
         cumulative_loss = 0.0
@@ -291,7 +311,10 @@ def evaluate(model, data, epoch, args, tb_writer=None):
                     # however, system RAM is easily exceeded and compute time becomes problematic
                     all_image_features.append(image_features.cpu())
                     all_text_features.append(text_features.cpu())
-                    logit_scale = logit_scale.mean()
+                    if not args.force_mrl_loss: 
+                        logit_scale = logit_scale.mean()  
+                    else:
+                        logit_scale = sum(logit_scale)/len(logit_scale)
                     logits_per_image = logit_scale * image_features @ text_features.t()
                     logits_per_text = logits_per_image.t()
 

From f3668c2b3d0a12a134022f334273cff0df6924d6 Mon Sep 17 00:00:00 2001
From: krmayankb <krmayankb@gmail.com>
Date: Thu, 9 Nov 2023 19:30:41 -0800
Subject: [PATCH 07/18] normalize with F

---
 src/open_clip/model.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/open_clip/model.py b/src/open_clip/model.py
index 9f7ce51..d2e3f2d 100644
--- a/src/open_clip/model.py
+++ b/src/open_clip/model.py
@@ -230,8 +230,8 @@ def xm_normalize(self, features, normalize: bool=False):
 
     def encode_image(self, image, normalize: bool = False):
         features = self.visual(image)
-        if xm.xla_device():
-            return self.xm_normalize(features, normalize=normalize)
+        # if xm.xla_device():
+        #     return self.xm_normalize(features, normalize=normalize)
         return F.normalize(features, dim=-1) if normalize else features
 
     def encode_text(self, text, normalize: bool = False):
@@ -246,8 +246,8 @@ def encode_text(self, text, normalize: bool = False):
         x = self.ln_final(x)  # [batch_size, n_ctx, transformer.width]
         # take features from the eot embedding (eot_token is the highest number in each sequence)
         x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
-        if xm.xla_device():
-            return self.xm_normalize(x, normalize=normalize)
+        # if xm.xla_device():
+        #     return self.xm_normalize(x, normalize=normalize)
         return F.normalize(x, dim=-1) if normalize else x
 
     def forward(self, image, text):

From 75924207a1ce207cf246245970f1c1e552e85fba Mon Sep 17 00:00:00 2001
From: krmayankb <krmayankb@gmail.com>
Date: Thu, 9 Nov 2023 19:41:10 -0800
Subject: [PATCH 08/18] autocast fix for tpu

---
 src/training/zero_shot.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/training/zero_shot.py b/src/training/zero_shot.py
index ac814c8..58afa53 100644
--- a/src/training/zero_shot.py
+++ b/src/training/zero_shot.py
@@ -3,6 +3,7 @@
 import torch
 import torch.nn.functional as F
 from tqdm import tqdm
+import contextlib
 
 from open_clip import get_cast_dtype, get_tokenizer
 from .precision import get_autocast
@@ -36,7 +37,7 @@ def accuracy(output, target, topk=(1,)):
 
 
 def run(model, classifier, dataloader, dim, args):
-    autocast = get_autocast(args.precision)
+    autocast = get_autocast(args.precision) if not args.use_tpu else contextlib.nullcontext
     cast_dtype = get_cast_dtype(args.precision)
     with torch.no_grad():
         top1, top5, n = 0., 0., 0.

From 575574d723cb1409679a182d7c83941de275482f Mon Sep 17 00:00:00 2001
From: krmayankb <krmayankb@gmail.com>
Date: Thu, 9 Nov 2023 20:55:40 -0800
Subject: [PATCH 09/18] change normalize to F

---
 src/open_clip/loss.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/open_clip/loss.py b/src/open_clip/loss.py
index 681f071..4580c12 100644
--- a/src/open_clip/loss.py
+++ b/src/open_clip/loss.py
@@ -252,8 +252,10 @@ def forward(self, image_features, text_features, logit_scale, output_dict=False)
         # total_loss = 0
         loss_list = []
         for idx, dim in enumerate(self.dim_to_consider): 
-            img = self.normalize(image_features[:,:dim], dim=-1) # slice and normalize 
-            txt = self.normalize(text_features[:,:dim], dim=-1) # slice and normalize 
+            # img = self.normalize(image_features[:,:dim], dim=-1) # slice and normalize 
+            # txt = self.normalize(text_features[:,:dim], dim=-1) # slice and normalize 
+            img = F.normalize(image_features[:,:dim], dim=-1) # slice and normalize 
+            txt = F.normalize(text_features[:,:dim], dim=-1) # slice and normalize 
             loss = super().forward(image_features=img, text_features=txt, logit_scale=logit_scale[idx])
             # total_loss += self.mrl_loss_weights[idx] * loss
             loss_list.append(self.mrl_loss_weights[idx] * loss)

From f854cea57a2c52c87fadb475be5e93634f060ddb Mon Sep 17 00:00:00 2001
From: krmayankb <krmayankb@gmail.com>
Date: Sat, 11 Nov 2023 00:23:34 -0800
Subject: [PATCH 10/18] parallel loading logic update

---
 src/training/main.py      | 6 ++++++
 src/training/zero_shot.py | 2 +-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/training/main.py b/src/training/main.py
index 388fe1c..c36c26a 100644
--- a/src/training/main.py
+++ b/src/training/main.py
@@ -358,6 +358,12 @@ def main(args):
         data = get_data(args, (preprocess_train, preprocess_val), epoch=start_epoch, tokenizer=get_tokenizer(args.model))      
         if xm.is_master_ordinal():
             xm.rendezvous('download_only_once')  
+        
+        for split in ["train", "val"]:
+            device = xm.xla_device()
+            para_loader = pl.ParallelLoader(data[split].dataloader, [device])
+            data[split].dataloader = para_loader.per_device_loader(device) 
+    
     else:
         data = get_data(args, (preprocess_train, preprocess_val), epoch=start_epoch, tokenizer=get_tokenizer(args.model))
     assert len(data), 'At least one train or eval dataset must be specified.'
diff --git a/src/training/zero_shot.py b/src/training/zero_shot.py
index 58afa53..245cd63 100644
--- a/src/training/zero_shot.py
+++ b/src/training/zero_shot.py
@@ -38,7 +38,7 @@ def accuracy(output, target, topk=(1,)):
 
 def run(model, classifier, dataloader, dim, args):
     autocast = get_autocast(args.precision) if not args.use_tpu else contextlib.nullcontext
-    cast_dtype = get_cast_dtype(args.precision)
+    cast_dtype = get_cast_dtype(args.precision)    
     with torch.no_grad():
         top1, top5, n = 0., 0., 0.
         for images, target in tqdm(dataloader, unit_scale=args.batch_size):

From c927df46efc6a809a0c626d87b0d90e11024128b Mon Sep 17 00:00:00 2001
From: krmayankb <krmayankb@gmail.com>
Date: Sat, 11 Nov 2023 00:32:23 -0800
Subject: [PATCH 11/18] split fix

---
 src/training/main.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/training/main.py b/src/training/main.py
index c36c26a..f1a4bff 100644
--- a/src/training/main.py
+++ b/src/training/main.py
@@ -359,7 +359,7 @@ def main(args):
         if xm.is_master_ordinal():
             xm.rendezvous('download_only_once')  
         
-        for split in ["train", "val"]:
+        for split in ["train", "imagenet-val"]:
             device = xm.xla_device()
             para_loader = pl.ParallelLoader(data[split].dataloader, [device])
             data[split].dataloader = para_loader.per_device_loader(device) 

From 490d7a5c14b1102c0a376859ee5962b13f7003e6 Mon Sep 17 00:00:00 2001
From: krmayankb <krmayankb@gmail.com>
Date: Sat, 11 Nov 2023 00:37:35 -0800
Subject: [PATCH 12/18] bug fix

---
 src/training/main.py | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/src/training/main.py b/src/training/main.py
index f1a4bff..31cbc3d 100644
--- a/src/training/main.py
+++ b/src/training/main.py
@@ -358,12 +358,6 @@ def main(args):
         data = get_data(args, (preprocess_train, preprocess_val), epoch=start_epoch, tokenizer=get_tokenizer(args.model))      
         if xm.is_master_ordinal():
             xm.rendezvous('download_only_once')  
-        
-        for split in ["train", "imagenet-val"]:
-            device = xm.xla_device()
-            para_loader = pl.ParallelLoader(data[split].dataloader, [device])
-            data[split].dataloader = para_loader.per_device_loader(device) 
-    
     else:
         data = get_data(args, (preprocess_train, preprocess_val), epoch=start_epoch, tokenizer=get_tokenizer(args.model))
     assert len(data), 'At least one train or eval dataset must be specified.'
@@ -421,6 +415,11 @@ def main(args):
         return
 
     loss = create_loss(args)
+    if args.use_tpu: 
+        for split in ["train", "imagenet-val"]:
+        device = xm.xla_device()
+        para_loader = pl.ParallelLoader(data[split].dataloader, [device])
+        data[split].dataloader = para_loader.per_device_loader(device) 
 
     for epoch in range(start_epoch, args.epochs):
         if is_master(args):

From b3a2a24dc516392c4d3b66b8e0b25486a93a34ba Mon Sep 17 00:00:00 2001
From: krmayankb <krmayankb@gmail.com>
Date: Sat, 11 Nov 2023 00:38:36 -0800
Subject: [PATCH 13/18] bug fix

---
 src/training/main.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/training/main.py b/src/training/main.py
index 31cbc3d..e5ff30e 100644
--- a/src/training/main.py
+++ b/src/training/main.py
@@ -417,9 +417,9 @@ def main(args):
     loss = create_loss(args)
     if args.use_tpu: 
         for split in ["train", "imagenet-val"]:
-        device = xm.xla_device()
-        para_loader = pl.ParallelLoader(data[split].dataloader, [device])
-        data[split].dataloader = para_loader.per_device_loader(device) 
+            device = xm.xla_device()
+            para_loader = pl.ParallelLoader(data[split].dataloader, [device])
+            data[split].dataloader = para_loader.per_device_loader(device) 
 
     for epoch in range(start_epoch, args.epochs):
         if is_master(args):

From b8bfd103203cac2661256e41f36e08e6d36aff47 Mon Sep 17 00:00:00 2001
From: krmayankb <krmayankb@gmail.com>
Date: Mon, 13 Nov 2023 22:05:27 -0800
Subject: [PATCH 14/18] refactor for xmp.spawn

---
 src/open_clip/model.py    | 11 -----------
 src/training/main.py      | 16 ++++++++--------
 src/training/zero_shot.py |  8 ++++++++
 3 files changed, 16 insertions(+), 19 deletions(-)

diff --git a/src/open_clip/model.py b/src/open_clip/model.py
index d2e3f2d..991fd60 100644
--- a/src/open_clip/model.py
+++ b/src/open_clip/model.py
@@ -221,17 +221,8 @@ def set_grad_checkpointing(self, enable=True):
         self.visual.set_grad_checkpointing(enable)
         self.transformer.grad_checkpointing = enable
 
-    def xm_normalize(self, features, normalize: bool=False):
-        if normalize:
-            norm = xm.all_reduce("sum", features ** 2)
-            norm = torch.sqrt(norm)
-            features = features / norm    
-        return features
-
     def encode_image(self, image, normalize: bool = False):
         features = self.visual(image)
-        # if xm.xla_device():
-        #     return self.xm_normalize(features, normalize=normalize)
         return F.normalize(features, dim=-1) if normalize else features
 
     def encode_text(self, text, normalize: bool = False):
@@ -246,8 +237,6 @@ def encode_text(self, text, normalize: bool = False):
         x = self.ln_final(x)  # [batch_size, n_ctx, transformer.width]
         # take features from the eot embedding (eot_token is the highest number in each sequence)
         x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
-        # if xm.xla_device():
-        #     return self.xm_normalize(x, normalize=normalize)
         return F.normalize(x, dim=-1) if normalize else x
 
     def forward(self, image, text):
diff --git a/src/training/main.py b/src/training/main.py
index e5ff30e..a49e2de 100644
--- a/src/training/main.py
+++ b/src/training/main.py
@@ -78,7 +78,7 @@ def get_latest_checkpoint(path: str, remote : bool):
     return None
 
 
-def main(args):
+def main(index, args):
     args = parse_args(args)
     # check if args.use_tpu is set, if so, use torch_xla        
     if torch.cuda.is_available():
@@ -415,12 +415,7 @@ def main(args):
         return
 
     loss = create_loss(args)
-    if args.use_tpu: 
-        for split in ["train", "imagenet-val"]:
-            device = xm.xla_device()
-            para_loader = pl.ParallelLoader(data[split].dataloader, [device])
-            data[split].dataloader = para_loader.per_device_loader(device) 
-
+    
     for epoch in range(start_epoch, args.epochs):
         if is_master(args):
             logging.info(f'Start epoch {epoch}')
@@ -496,8 +491,13 @@ def copy_codebase(args):
     return 1
 
 
+def run_tpu():
+    # Spawn 8 processes
+    xmp.spawn(main, args=(sys.argv[1:],))
+
+
 if __name__ == "__main__":
     if xmp is not None:
-        xmp.spawn(main(sys.argv[1:]), nprocs=8, start_method='fork')
+        run_tpu()
     else:
         main(sys.argv[1:])
diff --git a/src/training/zero_shot.py b/src/training/zero_shot.py
index 245cd63..fb88179 100644
--- a/src/training/zero_shot.py
+++ b/src/training/zero_shot.py
@@ -4,6 +4,8 @@
 import torch.nn.functional as F
 from tqdm import tqdm
 import contextlib
+import torch_xla.core.xla_model as xm
+import torch_xla.distributed.parallel_loader as pl
 
 from open_clip import get_cast_dtype, get_tokenizer
 from .precision import get_autocast
@@ -39,6 +41,12 @@ def accuracy(output, target, topk=(1,)):
 def run(model, classifier, dataloader, dim, args):
     autocast = get_autocast(args.precision) if not args.use_tpu else contextlib.nullcontext
     cast_dtype = get_cast_dtype(args.precision)    
+    
+    if args.use_tpu: 
+        device = xm.xla_device()
+        para_loader = pl.ParallelLoader(dataloader, [device])
+        dataloader = para_loader.per_device_loader(device) 
+    
     with torch.no_grad():
         top1, top5, n = 0., 0., 0.
         for images, target in tqdm(dataloader, unit_scale=args.batch_size):

From bccdfda9bd44e58300f53a779587511228cd4d6e Mon Sep 17 00:00:00 2001
From: krmayankb <krmayankb@gmail.com>
Date: Tue, 21 Nov 2023 13:22:03 -0800
Subject: [PATCH 15/18] only master logging

---
 src/training/distributed.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/training/distributed.py b/src/training/distributed.py
index c120f1b..5f44dd1 100644
--- a/src/training/distributed.py
+++ b/src/training/distributed.py
@@ -22,6 +22,8 @@ def is_local_master(args):
 
 
 def is_master(args, local=False):
+    if xm is not None:
+        return xm.is_master_ordinal()
     return is_local_master(args) if local else is_global_master(args)
 
 

From e902d2da08fab70267868e699c6a9efdf0d662c2 Mon Sep 17 00:00:00 2001
From: Mayank Kumar <krmayankb@gmail.com>
Date: Thu, 28 Dec 2023 04:14:38 +0000
Subject: [PATCH 16/18] tpu training scripts

---
 scripts/training_tpu.sh | 37 +++++++++++++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)
 create mode 100644 scripts/training_tpu.sh

diff --git a/scripts/training_tpu.sh b/scripts/training_tpu.sh
new file mode 100644
index 0000000..46a396b
--- /dev/null
+++ b/scripts/training_tpu.sh
@@ -0,0 +1,37 @@
+#!/bin/bash
+
+export PJRT_DEVICE=TPU
+export XLA_IR_DEBUG=1
+export XLA_METRICS_FILE=1
+#sudo kill -9 $(sudo lsof -w /dev/accel0 | awk 'NR>1 {print $2}' | uniq)
+
+python -c "import os; os.environ.pop('LD_PRELOAD', None)"
+
+
+cd ~/open_clip/src/ 
+
+python3 -m training.main \
+    --model "ViT-B-16" \
+    --train-data "/home/krmayank/data/medium/data/{00000000..00001919}.tar" \
+    --imagenet-val="/home/krmayank/data/imagenet/imagenet_val/" \
+    --precision "amp_bfloat16" \
+    --dataset-type webdataset \
+    --use_tpu \
+    --force_mrl_loss \
+    --mrl_dim_to_consider "768,384,192,96,48" \
+    --gather-with-grad \
+    --local-loss \
+    --batch-size 64 \
+    --accum-freq 1 \
+    --workers 2 \
+    --epochs 3 \
+    --warmup 1 \
+    --zeroshot-frequency 1 \
+    --seed 0 \
+    --logs "../logs/" \
+    --train-num-samples 10000 \
+    --val-num-samples 10000 \
+#    --report-to 'wandb' \
+#    --wandb-project-name "mrl_clip_training" 
+# --val-data "/home/krmayank/data/medium/data/{00000000..00000000}.tar" \
+# --imagenet-val="/home/krmayank/data/imagenet/imagenet_val/" \

From 3576ebb6d012e6dcafa16ee238e9d8b650a223c8 Mon Sep 17 00:00:00 2001
From: Mayank Kumar <krmayankb@gmail.com>
Date: Mon, 29 Jan 2024 00:15:20 +0000
Subject: [PATCH 17/18] adding mark_steps

---
 scripts/mrl_clip.sh       | 31 +++++++++++++++++--------------
 scripts/training_tpu.sh   | 23 +++++++++++------------
 src/training/zero_shot.py |  8 +++++++-
 3 files changed, 35 insertions(+), 27 deletions(-)

diff --git a/scripts/mrl_clip.sh b/scripts/mrl_clip.sh
index f991e61..45b9ac4 100644
--- a/scripts/mrl_clip.sh
+++ b/scripts/mrl_clip.sh
@@ -1,23 +1,26 @@
-cd /mmfs1/gscratch/krishna/mayank/clip_clone/open_clip/src
-torchrun --nproc_per_node 2 --master_port 3233 -m training.main \
-    --model "ViT-B-16" \
-    --train-data "/mmfs1/data/yfcc-tmp/cc_3m/train_shards/shard_{000000..003318}.tar" \
-    --imagenet-val "/mmfs1/data/yfcc-tmp/imagenet/val/" \
+cd ~/open_clip/src
+#torchrun --nproc_per_node 2 --master_port 3233 -m training.main \
+#lightning run model training/main.py \
+python3 -m training.main \
+    --model "ViT-B-32" \
+    --train-data "/home/krmayank/data/medium/data/{00000000..00001919}.tar" \
+    --imagenet-val="/home/krmayank/data/imagenet/imagenet_val/"
     --dataset-type webdataset \
     --precision amp \
     --gather-with-grad \
     --local-loss \
     --force_mrl_loss \
-    --mrl_loss_weights "0.1,0.15,0.2,0.25,0.3" \
+    --mrl_loss_weights "1,1,1,1,1" \
     --mrl_dim_to_consider "768,384,192,96,48" \
-    --batch-size 512 \
+    --batch-size 128 \
     --accum-freq 1 \
     --workers 4 \
-    --epochs 40 \
+    --epochs 3 \
     --warmup 4000 \
-    --zeroshot-frequency 2 \
-    --seed 0 \
-    --report-to 'wandb' \
-    --wandb-project-name "mrl_clip_training" \
-    --logs "/mmfs1/gscratch/krishna/mayank/clip_clone/open_clip/src/logs/mrl_clip" \
-    --name "mrl_clip_b512_accum_1_ep40_diffLogitScale_D082723_wl010150202503" 
\ No newline at end of file
+    --zeroshot-frequency 1 \
+    --seed 1234 \
+    --num_train_samples 10000 
+    #--report-to 'wandb' \
+    #--wandb-project-name "mrl_clip_training" \
+    #--logs "/mmfs1/gscratch/krishna/mayank/clip_clone/open_clip/src/logs/mrl_clip" \
+    #--name "mrl_clip_b512_accum_1_ep40_diffLogitScale_D082723_wl010150202503" 
diff --git a/scripts/training_tpu.sh b/scripts/training_tpu.sh
index 46a396b..d1f53d6 100644
--- a/scripts/training_tpu.sh
+++ b/scripts/training_tpu.sh
@@ -9,28 +9,27 @@ python -c "import os; os.environ.pop('LD_PRELOAD', None)"
 
 
 cd ~/open_clip/src/ 
-
 python3 -m training.main \
-    --model "ViT-B-16" \
+    --model "ViT-B-32" \
     --train-data "/home/krmayank/data/medium/data/{00000000..00001919}.tar" \
     --imagenet-val="/home/krmayank/data/imagenet/imagenet_val/" \
-    --precision "amp_bfloat16" \
     --dataset-type webdataset \
-    --use_tpu \
-    --force_mrl_loss \
-    --mrl_dim_to_consider "768,384,192,96,48" \
+    --precision amp \
     --gather-with-grad \
     --local-loss \
-    --batch-size 64 \
+    --force_mrl_loss \
+    --mrl_loss_weights "1,1,1,1,1" \
+    --mrl_dim_to_consider "768,384,192,96,48" \
+    --batch-size 128 \
     --accum-freq 1 \
-    --workers 2 \
+    --workers 4 \
     --epochs 3 \
-    --warmup 1 \
+    --warmup 4 \
     --zeroshot-frequency 1 \
-    --seed 0 \
-    --logs "../logs/" \
+    --seed 1234 \
+    --gather-with-grad \
     --train-num-samples 10000 \
-    --val-num-samples 10000 \
+    --val-num-samples 10000
 #    --report-to 'wandb' \
 #    --wandb-project-name "mrl_clip_training" 
 # --val-data "/home/krmayank/data/medium/data/{00000000..00000000}.tar" \
diff --git a/src/training/zero_shot.py b/src/training/zero_shot.py
index fb88179..b3731e3 100644
--- a/src/training/zero_shot.py
+++ b/src/training/zero_shot.py
@@ -13,7 +13,9 @@
 
 
 def zero_shot_classifier(model, classnames, templates, dim, args):
+    xm.mark_step()
     tokenizer = get_tokenizer(args.model)
+    xm.mark_step()
     with torch.no_grad():
         zeroshot_weights = []
         for classname in tqdm(classnames):
@@ -28,7 +30,9 @@ def zero_shot_classifier(model, classnames, templates, dim, args):
             class_embedding = F.normalize(class_embeddings, dim=-1).mean(dim=0)
             class_embedding /= class_embedding.norm()
             zeroshot_weights.append(class_embedding)
-        zeroshot_weights = torch.stack(zeroshot_weights, dim=1).to(args.device)
+            xm.mark_step()
+       zeroshot_weights = torch.stack(zeroshot_weights, dim=1).to(args.device)
+       xm.mark_step()
     return zeroshot_weights
 
 
@@ -67,7 +71,9 @@ def run(model, classifier, dataloader, dim, args):
                 logits = 100. * image_features @ classifier
 
             # measure accuracy
+            xm.mark_step()
             acc1, acc5 = accuracy(logits, target, topk=(1, 5))
+            xm.mark_step()
             top1 += acc1
             top5 += acc5
             n += images.size(0)

From 8cb00d468c145cbcda172dd270f9dd580e26892f Mon Sep 17 00:00:00 2001
From: Mayank Kumar <krmayankb@gmail.com>
Date: Mon, 29 Jan 2024 00:35:48 +0000
Subject: [PATCH 18/18] indent fix

---
 scripts/training_tpu.sh   | 5 ++---
 src/training/zero_shot.py | 4 ++--
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/scripts/training_tpu.sh b/scripts/training_tpu.sh
index d1f53d6..415354a 100644
--- a/scripts/training_tpu.sh
+++ b/scripts/training_tpu.sh
@@ -1,13 +1,12 @@
 #!/bin/bash
 
 export PJRT_DEVICE=TPU
-export XLA_IR_DEBUG=1
-export XLA_METRICS_FILE=1
+#export XLA_IR_DEBUG=1
+#export XLA_METRICS_FILE=1
 #sudo kill -9 $(sudo lsof -w /dev/accel0 | awk 'NR>1 {print $2}' | uniq)
 
 python -c "import os; os.environ.pop('LD_PRELOAD', None)"
 
-
 cd ~/open_clip/src/ 
 python3 -m training.main \
     --model "ViT-B-32" \
diff --git a/src/training/zero_shot.py b/src/training/zero_shot.py
index b3731e3..9f062a1 100644
--- a/src/training/zero_shot.py
+++ b/src/training/zero_shot.py
@@ -31,8 +31,8 @@ def zero_shot_classifier(model, classnames, templates, dim, args):
             class_embedding /= class_embedding.norm()
             zeroshot_weights.append(class_embedding)
             xm.mark_step()
-       zeroshot_weights = torch.stack(zeroshot_weights, dim=1).to(args.device)
-       xm.mark_step()
+        zeroshot_weights = torch.stack(zeroshot_weights, dim=1).to(args.device)
+        xm.mark_step()
     return zeroshot_weights