linear callback rank

Author: karinazad

src/lobster/callbacks/_linear_probe_callback.py

@@ -73,6 +73,8 @@ def _skip(self, trainer: L.Trainer) -> bool:
         # Don't skip if global_rank=0 (main process)
         if trainer.global_rank == 0:
             return False
+        else:
+            return True
 
         return trainer.current_epoch % self.run_every_n_epochs != 0
 

src/lobster/hydra_config/experiment/train_ume.yaml

@@ -1,38 +1,42 @@
 # @package _global_
 
-# lobster_train experiment=train_ume logger.entity=zadorozk logger.project=lobster logger.group=ume-dev ++trainer.accelerator=gpu ++trainer.devices=auto 
+# lobster_train experiment=train_ume ++trainer.devices=auto 
 
 defaults:
   - override /model: modern_bert.yaml
   - override /data: ume
-  - override /callbacks: [base, tokens_per_second, moleculeace_linear_probe_fast, calm_linear_probe_fast]
+  - override /callbacks: [base, moleculeace_linear_probe_fast, throughput, tokens_per_second]
+
+compile: true
 
 data:
   _target_: lobster.data.UmeLightningDataModule
   root: ${paths.root_dir}/data
   datasets: ["M320M", "Calm", "AMPLIFY"]
-  batch_size: 16
+  batch_size: 128
   tokenizer_max_length: ${model.max_length}
   pin_memory: true
-  shuffle_buffer_size: 10_000
-  num_workers: 32
+  shuffle_buffer_size: 1000
+  num_workers: 8
   seed: 0
 
 paths:
-  root_dir: dev
+  root_dir: ./runs
 
 trainer:
   max_steps: 50_000
-  val_check_interval: 1_000
+  val_check_interval: 2000
   precision: 16-mixed
   gradient_clip_val: 0.5
   accumulate_grad_batches: 8
   max_time: "00:24:00:00"
-
+  limit_val_batches: 20_000
+  #limit_train_batches: 1000 # for debugging
+  #profiler: simple
 
 model:
-  model_name: UME_mini
-  vocab_size:  627
+  model_name: UME_medium
+  vocab_size:  640
   pad_token_id: 1
   cls_token_id: 0
   mask_token_id: 4
@@ -51,10 +55,8 @@ model:
 callbacks:
   moleculeace_linear_probe:
     max_length: ${model.max_length}
-  calm_linear_probe:
-    max_length: ${model.max_length}
 
 logger:
   name: ume_amplify_m320m_calm_${model.model_name}_${now:%Y-%m-%d_%H-%M-%S}
   project: lobster
-  group: ume-dev
+  group: ume-dev-${now:%Y-%m-%d-%H-%M-%S}