Validation step3

examples/config_validation_tiny_llama.yaml

@@ -0,0 +1,134 @@
+checkpoints:
+  checkpoint_interval: 10
+  checkpoints_path: checkpoints
+  checkpoints_path_is_shared_file_system: false
+  resume_checkpoint_path: null
+  save_initial_state: false
+data_stages:
+- data:
+    dataset:
+      dataset_overwrite_cache: false
+      dataset_processing_num_proc_per_process: 1
+      hf_dataset_config_name: null
+      hf_dataset_or_datasets: stas/openwebtext-10k
+      hf_dataset_splits: train
+      text_column_name: text
+    num_loading_workers: 1
+    seed: 42
+  name: Stable Training Stage
+  start_training_step: 1
+- data:
+    dataset:
+      dataset_overwrite_cache: false
+      dataset_processing_num_proc_per_process: 1
+      hf_dataset_config_name: null
+      hf_dataset_or_datasets: stas/openwebtext-10k
+      hf_dataset_splits: train
+      text_column_name: text
+    num_loading_workers: 1
+    seed: 42
+  name: Annealing Phase
+  start_training_step: 10
+valid_data_stages:
+- data:
+    dataset:
+      dataset_overwrite_cache: false
+      dataset_processing_num_proc_per_process: 1
+      hf_dataset_config_name: null
+      hf_dataset_or_datasets: stas/oscar-en-10k
+      hf_dataset_splits: train
+      text_column_name: text
+    num_loading_workers: 1
+    seed: 42
+  name: Stable Training Stage
+  start_training_step: 1
+- data:
+    dataset:
+      dataset_overwrite_cache: false
+      dataset_processing_num_proc_per_process: 1
+      hf_dataset_config_name: null
+      hf_dataset_or_datasets: stas/oscar-en-10k
+      hf_dataset_splits: train
+      text_column_name: text
+    num_loading_workers: 1
+    seed: 42
+  name: Annealing Phase
+  start_training_step: 8
+general:
+  benchmark_csv_path: null
+  consumed_train_samples: null
+  ignore_sanity_checks: true
+  project: debug
+  run: tiny_llama_%date_%jobid
+  seed: 42
+  step: null
+lighteval: null
+logging:
+  iteration_step_info_interval: 1
+  log_level: info
+  log_level_replica: info
+model:
+  ddp_bucket_cap_mb: 25
+  dtype: float32
+  init_method:
+    std: 0.025
+  make_vocab_size_divisible_by: 1
+  model_config:
+    bos_token_id: 1
+    eos_token_id: 2
+    hidden_act: silu
+    hidden_size: 16
+    initializer_range: 0.02
+    intermediate_size: 64
+    is_llama_config: true
+    max_position_embeddings: 256
+    num_attention_heads: 4
+    num_hidden_layers: 2
+    num_key_value_heads: 4
+    pad_token_id: null
+    pretraining_tp: 1
+    rms_norm_eps: 1.0e-05
+    rope_scaling: null
+    tie_word_embeddings: true
+    use_cache: true
+    vocab_size: 256
+optimizer:
+  accumulate_grad_in_fp32: true
+  clip_grad: 1.0
+  learning_rate_scheduler:
+    learning_rate: 0.0003
+    lr_decay_starting_step: null
+    lr_decay_steps: 13
+    lr_decay_style: cosine
+    lr_warmup_steps: 2
+    lr_warmup_style: linear
+    min_decay_lr: 1.0e-05
+  optimizer_factory:
+    adam_beta1: 0.9
+    adam_beta2: 0.95
+    adam_eps: 1.0e-08
+    name: adamW
+    torch_adam_is_fused: true
+  weight_decay: 0.01
+  zero_stage: 0
+parallelism:
+  dp: 1
+  expert_parallel_size: 1
+  pp: 1
+  pp_engine: 1f1b
+  tp: 1
+  tp_linear_async_communication: true
+  tp_mode: REDUCE_SCATTER
+profiler: null
+tokenizer:
+  tokenizer_max_length: null
+  tokenizer_name_or_path: robot-test/dummy-tokenizer-wordlevel
+  tokenizer_revision: null
+tokens:
+  batch_accumulation_per_replica: 1
+  limit_test_batches: 0
+  limit_val_batches: 5
+  micro_batch_size: 2
+  sequence_length: 256
+  train_steps: 200
+  val_check_interval: 2

reconcile_rotary_embeddings.py

@@ -0,0 +1,91 @@
+import torch
+import torchtune
+import flash_attn
+import flash_attn.layers.rotary
+
+
+class RotaryEmbeddingKyleLikeFA(torch.nn.Module):
+    """
+    Has the same function signature as FA, for interleaved=True and separate q, kv. 
+    seqlen_offset = 0
+    Does not operate inplace, but that's fine for how it's used in Nanotron. 
+    """
+    def __init__(self, dim: int, base: float):
+        super().__init__()
+        self.dim = dim
+        self.base = float(base)
+
+        self.max_seq_len = None
+        self.rpe = None
+
+    def forward(self, q, kv):
+        bs, q_len, n_heads, _ = q.shape
+        assert self.dim == _
+
+        assert (bs, q_len, 2, n_heads, self.dim) == kv.shape
+
+        if (self.rpe is None) or (self.max_seq_len != q_len):
+            self.max_seq_len = q_len 
+            self.rpe = torchtune.modules.RotaryPositionalEmbeddings(dim=self.dim, 
+                                                                    max_seq_len=self.max_seq_len,
+                                                                    base=self.base).to(device)
+        q_out = self.rpe(q)
+        kv_out = torch.stack((self.rpe(kv[:, :, 0]), kv[:, :, 1]), 2)
+        return q_out, kv_out
+
+
+
+if __name__ == "__main__":
+    device = torch.device(0) 
+    theta = 10000
+
+    batch_size = 3
+    dim_qk = 4
+    q_len = 256
+    kv_len = 256
+    n_heads = 4
+
+    max_seq_len = max(q_len, kv_len) 
+
+    print(max_seq_len) 
+
+
+    query_states = torch.rand(batch_size, q_len, n_heads, dim_qk, device=device) 
+    key_value_states = torch.rand(batch_size, kv_len, 2, n_heads, dim_qk, device=device).contiguous()
+
+
+    interleaved = True 
+    # interleaved = False
+    re1 = flash_attn.layers.rotary.RotaryEmbedding(dim=dim_qk, interleaved=interleaved, base=theta).to(device)
+    re2 = torchtune.modules.RotaryPositionalEmbeddings(dim=dim_qk, max_seq_len=max_seq_len, base=theta).to(device)
+    re3 = RotaryEmbeddingKyleLikeFA(dim=dim_qk, base=theta).to(device)
+
+
+
+    print(key_value_states[:, :, 0].shape)
+
+    out2 = re2(query_states)
+    out3 = re2(key_value_states[:, :, 0]) 
+    # out4 = re2(key_value_states[:, :, 1]) 
+
+    out_eq = re3(query_states, kv=key_value_states)
+
+    # torch.testing.assert_close(out2, query_states)
+    out1 = re1(query_states, kv=key_value_states)
+
+    torch.testing.assert_close(out_eq[0], out1[0])
+    torch.testing.assert_close(out_eq[1], out1[1])
+
+
+    # Do this second, since the computation is inplace
+    torch.testing.assert_close(out1[0], query_states)
+
+    test = torch.stack((out3, key_value_states[:, :, 1]), 2)
+    torch.testing.assert_close(out1[1], test)
+    # torch.testing.assert_close(out1[1][:, :, 0], out3) 
+
+
+    torch.testing.assert_close(out1[0], out2)
+
+    print("done")
+

run_train.py

@@ -178,44 +178,48 @@ def get_dataloader_from_data_stage(
     return dataloader
 
 
-def get_dataloader(trainer: DistributedTrainer) -> Dict[str, DataLoader]:
+def get_dataloader(trainer: DistributedTrainer,
+                   data_stages_fieldname: str,
+                   metadata_fieldname: str) -> Dict[str, DataLoader]:
     dataloaders = {}
+    data_stages = getattr(trainer.config, data_stages_fieldname)
+    if data_stages:
+        metadata = getattr(trainer, metadata_fieldname)
+        for stage_idx, stage in enumerate(data_stages):
+            # NOTE: we only create the dataloader for the first stage,
+            # then we lazy initialize the dataloader for the other stages
+            stage = cast(DatasetStageArgs, stage)
+            consumed_train_samples = get_consumed_train_samples_of_a_data_stage_from_ckp(stage, metadata)
+            assert (
+                consumed_train_samples is not None
+            ), f"Cannot find consumed_train_samples for stage {stage.start_training_step} in the checkpoint"
 
-    for stage_idx, stage in enumerate(trainer.config.data_stages):
-        # NOTE: we only create the dataloader for the first stage,
-        # then we lazy initialize the dataloader for the other stages
-        stage = cast(DatasetStageArgs, stage)
-        consumed_train_samples = get_consumed_train_samples_of_a_data_stage_from_ckp(stage, trainer.metadata)
-        assert (
-            consumed_train_samples is not None
-        ), f"Cannot find consumed_train_samples for stage {stage.start_training_step} in the checkpoint"
-
-        num_remaining_train_steps = compute_remain_train_steps_of_a_data_stage_from_ckp(
-            stage, trainer.config, trainer.metadata
-        )
-        log_rank(
-            f"[Training Plan] Stage {stage.name} has {num_remaining_train_steps} remaining training steps and has consumed {consumed_train_samples} samples",
-            logger=logger,
-            level=logging.INFO,
-            rank=0,
-        )
-
-        dataloader = (
-            get_dataloader_from_data_stage(
-                trainer,
-                stage.data,
-                consumed_train_samples=consumed_train_samples,
-                num_remaining_train_steps=num_remaining_train_steps,
+            num_remaining_train_steps = compute_remain_train_steps_of_a_data_stage_from_ckp(
+                stage, trainer.config, metadata
             )
-            if stage_idx == 0
-            else lambda stage=stage: get_dataloader_from_data_stage(
-                trainer,
-                stage.data,
-                consumed_train_samples=consumed_train_samples,
-                num_remaining_train_steps=num_remaining_train_steps,
+            log_rank(
+                f"[Training Plan] Stage {stage.name} has {num_remaining_train_steps} remaining training steps and has consumed {consumed_train_samples} samples",
+                logger=logger,
+                level=logging.INFO,
+                rank=0,
             )
-        )
-        dataloaders[stage.name] = dataloader
+
+            dataloader = (
+                get_dataloader_from_data_stage(
+                    trainer,
+                    stage.data,
+                    consumed_train_samples=consumed_train_samples,
+                    num_remaining_train_steps=num_remaining_train_steps,
+                )
+                if stage_idx == 0
+                else lambda stage=stage: get_dataloader_from_data_stage(
+                    trainer,
+                    stage.data,
+                    consumed_train_samples=consumed_train_samples,
+                    num_remaining_train_steps=num_remaining_train_steps,
+                )
+            )
+            dataloaders[stage.name] = dataloader
     return dataloaders
 
 
@@ -231,7 +235,8 @@ def get_args():
 
     # Load trainer and data
     trainer = DistributedTrainer(config_file)
-    dataloader = get_dataloader(trainer)
+    dataloader_train = get_dataloader(trainer, "data_stages", "metadata")
+    dataloader_valid = get_dataloader(trainer, "valid_data_stages", "valid_metadata")
 
     # Train
-    trainer.train(dataloader)
+    trainer.train(dataloader_train, dataloader_valid)

src/nanotron/config/config.py

@@ -338,6 +338,7 @@ class Config:
     tokens: Optional[TokensArgs] = None
     optimizer: Optional[OptimizerArgs] = None
     data_stages: Optional[List[DatasetStageArgs]] = None
+    valid_data_stages: Optional[List[DatasetStageArgs]] = None
     profiler: Optional[ProfilerArgs] = None
     lighteval: Optional[LightEvalConfig] = None