Warmstart infrastructure switch

.gitignore

@@ -160,5 +160,5 @@ pyenv*
 noteboks/*
 
 tests/tmp/*
+*wandb_storage*
 .coverage/*
-wandb_storage/

config_files/training/config_example_coca.yaml

@@ -4,27 +4,53 @@ settings:
   referencing_keys:
     sample_key: input_ids
     target_key: target_ids
-  training:
-    training_log_interval_in_steps: 2
-    checkpointing_interval_in_steps: 2
-    evaluation_interval_in_steps: 2
-    global_num_seen_tokens: 0
-    activation_checkpointing_modules: []
-    gradient_acc_steps: 1
-    local_train_micro_batch_size: 3
-    sequence_length: 256
+    prediction_key: logits
   cuda_env:
     local_rank: ${cuda_env:LOCAL_RANK}
     global_rank: ${cuda_env:RANK}
     world_size: ${cuda_env:WORLD_SIZE}
   paths:
-    checkpointing_path: data/checkpoints
-
-tokenizer:
-  component_key: tokenizer
-  variant_key: gpt2_tokenizer_fast
-  config:
-    tokenizer_file: data/tokenizer/tokenizer_gpt2.json
+    checkpoint_saving_path: data/checkpoints
+    train_dataset_path: ./data/lorem_ipsum.pbin
+  intervals:
+    training_log_interval_in_steps: 2
+    checkpointing_interval_in_steps: 2
+    evaluation_interval_in_steps: 2
+  consistency_enforcement:
+    enforce_tokens_per_step_consistency: true
+    enforce_last_step_logged: false
+    enforce_last_step_evaluated: false
+    enforce_last_step_checkpointed: false
+  step_profile: 
+    gradient_accumulation_steps: 1
+    local_train_micro_batch_size: 1
+    sequence_length: 256
+  training_target:
+    num_target_tokens:      
+      component_key: number_conversion
+      variant_key: num_tokens_from_num_steps
+      config:
+        num_steps: ${settings.training_target.num_target_steps}
+        num_ranks: ${settings.cuda_env.world_size}
+        local_micro_batch_size: ${settings.step_profile.local_train_micro_batch_size}
+        sequence_length: ${settings.step_profile.sequence_length}
+        gradient_accumulation_steps: ${settings.step_profile.gradient_accumulation_steps}
+    num_target_steps:  # for the batch progress subscriber
+      component_key: number_conversion
+      variant_key: num_steps_from_num_samples
+      config:
+        num_ranks: ${settings.cuda_env.world_size}
+        local_micro_batch_size: ${settings.step_profile.local_train_micro_batch_size}
+        global_num_samples: ${settings.coca_example_settings.train_num_samples}
+        gradient_accumulation_steps: ${settings.step_profile.gradient_accumulation_steps}
+  training_progress: 
+    global_num_seen_tokens: 0
+    num_seen_steps: 0
+    local_num_seen_batches: 0
+    last_step: -1
+  coca_example_settings:
+    train_num_samples: 64
+    val_num_samples: 32
 
 collate_fn:
   component_key: collate_fn
@@ -41,7 +67,7 @@ train_dataset:
   component_key: dataset
   variant_key: dummy_dataset
   config:
-    num_samples: 64
+    num_samples: ${settings.coca_example_settings.train_num_samples}
     sample_definition:
       - sample_key: images
         sample_shape: [3, 224, 224]
@@ -54,7 +80,7 @@ val_dataset:
   component_key: dataset
   variant_key: dummy_dataset
   config:
-    num_samples: 32
+    num_samples: ${settings.coca_example_settings.val_num_samples}
     sample_definition:
       - sample_key: images
         sample_shape: [3, 224, 224]
@@ -69,23 +95,26 @@ train_dataloader:
   config:
     num_workers: 2
     pin_memory: true
-    shuffle: false
-    dataloader_tag: "train"
+    dataloader_tag: train
+    skip_num_batches: ${settings.training_progress.local_num_seen_batches}
     dataset:
       instance_key: train_dataset
       pass_type: BY_REFERENCE
     batch_sampler:
       component_key: batch_sampler
       variant_key: default
       config:
-        batch_size: ${settings.training.local_train_micro_batch_size}
+        batch_size: ${settings.step_profile.local_train_micro_batch_size}
+        drop_last: true
         sampler:
           component_key: sampler
           variant_key: distributed_sampler
           config:
             rank: ${settings.cuda_env.global_rank}
             num_replicas: ${settings.cuda_env.world_size}
             shuffle: true
+            drop_last: true
+            seed: 42
             dataset:
               instance_key: train_dataset
               pass_type: BY_REFERENCE
@@ -99,23 +128,25 @@ val_dataloader:
   config:
     num_workers: 2
     pin_memory: true
-    shuffle: false
-    dataloader_tag: "val"
+    dataloader_tag: val
     dataset:
       instance_key: val_dataset
       pass_type: BY_REFERENCE
     batch_sampler:
       component_key: batch_sampler
       variant_key: default
       config:
-        batch_size: ${settings.training.local_train_micro_batch_size}
+        batch_size: ${settings.step_profile.local_train_micro_batch_size}
+        drop_last: true
+
         sampler:
           component_key: sampler
           variant_key: distributed_sampler
           config:
             rank: ${settings.cuda_env.global_rank}
             num_replicas: ${settings.cuda_env.world_size}
             shuffle: false
+            drop_last: true
             dataset:
               instance_key: train_dataset
               pass_type: BY_REFERENCE
@@ -140,22 +171,16 @@ checkpoint_saving:
       component_key: checkpoint_saving_execution
       variant_key: fsdp
       config:
-        checkpoint_path: ${settings.paths.checkpointing_path}
+        checkpoint_path: ${settings.paths.checkpoint_saving_path}
         global_rank: ${settings.cuda_env.global_rank}
         experiment_id: ${settings.experiment_id}
-        get_num_tokens_from_num_steps_callable:
-          component_key: number_conversion
-          variant_key: num_tokens_from_num_steps_callable
-          config:
-            num_ranks: ${settings.cuda_env.world_size}
-            local_micro_batch_size: ${settings.training.local_train_micro_batch_size}
-            sequence_length: ${settings.training.sequence_length}
+
 loss_fn:
   component_key: loss
   variant_key: clm_cross_entropy_loss
   config:
     target_key: ${settings.referencing_keys.target_key}
-    prediction_key: logits
+    prediction_key: ${settings.referencing_keys.prediction_key}
 
 wrapped_model:
   component_key: model
@@ -169,7 +194,7 @@ wrapped_model:
     sharding_strategy: FULL_SHARD
     block_names: [TransformerBlock, VisionTransformerBlock]
 
-model: 
+model:
   component_key: model
   variant_key: model_initialized
   config:
@@ -241,9 +266,10 @@ scheduler:
     max_lr: 6e-4
     div_factor: 10
     final_div_factor: 1
-    total_steps: 64
+    total_steps: ${settings.training_target.num_target_steps}
     pct_start: 0.01
     anneal_strategy: cos
+    last_epoch: ${settings.training_progress.last_step}
 
 optimizer:
   component_key: optimizer
@@ -267,24 +293,14 @@ gradient_clipper:
       pass_type: BY_REFERENCE
     norm_type: P2_NORM
 
-
-batch_progress_subscriber:
+progress_subscriber:
   component_key: progress_subscriber
   variant_key: rich
   config:
     global_rank: ${settings.cuda_env.global_rank}
-    global_num_seen_steps:
-      component_key: number_conversion
-      variant_key: num_steps_from_num_tokens
-      config:
-        num_ranks: ${settings.cuda_env.world_size}
-        local_micro_batch_size: ${settings.training.local_train_micro_batch_size}
-        global_num_tokens: ${settings.training.global_num_seen_tokens}
-        sequence_length: ${settings.training.sequence_length}
-    gradient_acc_steps: ${settings.training.gradient_acc_steps}
-    train_dataloader:
-      instance_key: train_dataloader
-      pass_type: BY_REFERENCE
+    num_seen_steps: ${settings.training_progress.num_seen_steps}
+    num_target_steps: ${settings.training_target.num_target_steps}
+    train_dataloader_tag: ${train_dataloader.config.dataloader_tag}
     eval_dataloaders:
       instance_key: eval_dataloaders
       pass_type: BY_REFERENCE