WIP Kron

README.md

@@ -104,7 +104,7 @@ If you're using a TPU, more complete documentation for setting that up is availa
 As a kind of hello world, here's how you can train a GPT-2 "nano"-sized model on a small dataset.
 
 ```bash
-python -m levanter.main.train_lm --config_path config/gpt2_nano.yaml
+python -m levanter.main.train_lm --config_path config/llama2_100M_kron_test.yaml
 
 # alternatively, if you didn't use -e and are in a different directory
 python -m levanter.main.train_lm --config_path gpt2_nano

config/llama2_100M_kron_test.yaml

@@ -0,0 +1,34 @@
+data:
+  id: openwebtext
+model:
+  type: llama
+  seq_len: 4096
+  hidden_dim: 768
+  intermediate_dim: 3072
+  num_layers: 12
+  num_heads: 12
+  num_kv_heads: 12
+trainer:
+  tracker:
+    project: "levanter"
+    tags: ["pile", "llama"]
+  mp: p=f32,c=bfloat16
+  model_axis_size: 1
+  checkpointer:
+    keep:
+      - every: 1000
+    save_interval: 30m
+
+
+  train_batch_size: 1024
+  per_device_parallelism: 32  # set for v3 TPU
+  per_device_eval_parallelism: 32 # set a larger batch size for eval
+  num_train_steps: 50001
+optimizer:
+  learning_rate: 3E-4
+  weight_decay: 0.1
+  warmup: 2000
+  cooldown: 0.1
+  lr_schedule: constant
+  min_lr_ratio: 0.0 
+  type: kron

config/llama2_100M_mars.yaml

@@ -0,0 +1,34 @@
+data: !include data/dclm_gpt_neo.yaml
+model:
+  type: llama
+  seq_len: 4096
+  hidden_dim: 768
+  intermediate_dim: 3072
+  num_layers: 12
+  num_heads: 12
+  num_kv_heads: 12
+trainer:
+  tracker:
+    project: "levanter"
+    tags: ["pile", "llama"]
+  mp: p=f32,c=bfloat16
+  model_axis_size: 1
+  checkpointer:
+    keep:
+      - every: 1000
+    save_interval: 30m
+
+
+  train_batch_size: 1024
+  per_device_parallelism: 4  # set for v3 TPU
+  per_device_eval_parallelism: 4 # set a larger batch size for eval
+  num_train_steps: 50001
+optimizer:
+  learning_rate: 4E-3 # set low for fine-tuning
+  weight_decay: 0.1
+  min_lr_ratio: 0.0
+  warmup: 2000
+  cooldown: 0.4
+  lr_schedule: constant
+  gamma: 0.025
+  type: mars

config/llama2_100M_muon.yaml

@@ -0,0 +1,34 @@
+data: !include data/dclm_gpt_neo.yaml
+model:
+  type: llama
+  seq_len: 4096
+  hidden_dim: 768
+  intermediate_dim: 3072
+  num_layers: 12
+  num_heads: 12
+  num_kv_heads: 12
+trainer:
+  tracker:
+    project: "levanter"
+    tags: ["pile", "llama"]
+  mp: p=f32,c=bfloat16
+  model_axis_size: 1
+  checkpointer:
+    keep:
+      - every: 1000
+    save_interval: 30m
+
+
+  train_batch_size: 1024
+  per_device_parallelism: 4  # set for v3 TPU
+  per_device_eval_parallelism: 4 # set a larger batch size for eval
+  num_train_steps: 50001
+optimizer:
+  learning_rate: 2E-2 # set low for fine-tuning
+  weight_decay: 0
+  warmup: 0
+  cooldown: 0.1
+  lr_schedule: constant
+  min_lr_ratio: 0.0 
+  max_grad_norm: 0.0 
+  type: muon 

src/levanter/data/audio.py

@@ -193,7 +193,7 @@ def decode(x):
 
     def doc_iterator(self, split: str) -> Iterator[Tuple[np.ndarray, int, str]]:
         if self.id is not None:
-            data = datasets.load_dataset(self.id, split=split, name=self.name, streaming=self.stream)
+            data = datasets.load_dataset(self.id, split=split, name=self.name, streaming=self.stream, trust_remote_code=True)
             for doc in data:
                 yield (doc[self.audio_key]["array"], doc[self.audio_key]["sampling_rate"], doc[self.text_key])
         else:
@@ -385,7 +385,7 @@ def _has_validation_set(self):
 
         if self.id is not None:
             dataset = datasets.load_dataset(
-                self.id, name=self.name, streaming=self.stream, split=self.validation_split
+                self.id, name=self.name, streaming=self.stream, split=self.validation_split, trust_remote_code=True
             )
             try:
                 next(iter(dataset))

src/levanter/data/sharded_datasource.py

@@ -253,7 +253,7 @@ def open_shard_at_row(self, shard_name: str, row: int) -> Iterator[dict]:
     def _load_dataset(self):
         # obnoxiously, the dataset loading stuff doesn't work with ray because of multiprocessing
         # so we have to do this hacky thing where we load the dataset in the worker
-        return datasets.load_dataset(self.id, split=self.split, streaming=self.streaming, **self.kwargs)
+        return datasets.load_dataset(self.id, split=self.split, streaming=self.streaming, trust_remote_code=True, **self.kwargs)
 
 
 class TextUrlDataSource(ShardedDataSource[str]):

src/levanter/data/text.py

@@ -594,7 +594,7 @@ def get_shard_source(self, split) -> Optional[ShardedDataSource[str]]:
 
     def doc_iterator(self, split: str):
         if self.id is not None:
-            dataset = datasets.load_dataset(self.id, name=self.name, streaming=self.stream)
+            dataset = datasets.load_dataset(self.id, name=self.name, streaming=self.stream, trust_remote_code=True)
             data = dataset[split]
             for doc in data:
                 yield doc[self.text_key]
@@ -1065,7 +1065,7 @@ def _has_validation_set(self):
             return True
 
         if self.id is not None:
-            dataset = datasets.load_dataset(self.id, name=self.name, streaming=self.stream, split="validation")
+            dataset = datasets.load_dataset(self.id, name=self.name, streaming=self.stream, split="validation", trust_remote_code=True)
             try:
                 next(iter(dataset))
                 return True

src/levanter/optim/__init__.py

@@ -5,3 +5,12 @@
     scale_by_sophia_g,
     scale_by_sophia_h,
 )
+from .muon import (
+    MuonConfig,
+    ScaleByMuonState
+)
+from .mars import (
+    MarsConfig,
+    ScaleByMarsState
+)
+from .kron import KronConfig