wip: try make config run with python -m levanter.main.cache_dataset --config config/llama3_ul2r.yaml

Author: jyc

config/llama3_ul2r.yaml

@@ -0,0 +1,61 @@
+# Based on llama2_7b_continued.yaml.
+# TODO right now this is a `RayCachedLMDatasetConfig` for `cache_dataset.py`,
+# not `TrainLmConfig` for `train_lm.py`; change it to the latter.
+
+# data:
+id: dlwh/wikitext_103_detokenized
+tokenizer: meta-llama/Llama-3.1-8B
+format:
+  type: ul2r
+  text_key: text
+  task_configs:
+    r:
+      mask_prob: 0.15
+      mean_span_length: 3.0
+      random_roll: true
+      task_token: "<|reserved_special_token_3|>"
+    x1:
+      mask_prob: 0.15
+      mean_span_length: 32.0
+      random_roll: true
+      task_token: "<|reserved_special_token_4|>"
+    x2:
+      mask_prob: 0.5
+      mean_span_length: 3.0
+      random_roll: true
+      task_token: "<|reserved_special_token_4|>"
+    s:
+      task_token: "<|reserved_special_token_5|>"
+  task_probs:
+    r: 0.5
+    x1: 0.125
+    x2: 0.125
+    s: 0.25
+  rng_seed: 42
+
+# TODO haven't tested any of the model stuff yet
+
+# model:
+#   type: llama
+# initialize_from_hf: true
+# use_hf_model_config: true
+# model_name_or_path: meta-llama/Llama-3.1-8B
+
+# trainer:
+#   tracker:
+#     type: wandb
+#     project: "levanter"
+#     tags: ["ul2r", "llama3", "wikitext"]
+
+#   mp: p=f32,c=bfloat16
+
+#   model_axis_size: 1
+#   per_device_eval_parallelism: 4
+
+#   train_batch_size: 1024
+#   num_train_steps: 10000
+#   steps_per_eval: 500
+
+# optimizer:
+#   learning_rate: 1.2e-4
+#   weight_decay: 0.0

src/levanter/data/text.py

@@ -74,6 +74,7 @@
     UrlDataSource,
     WrappedHFDataSource,
 )
+from levanter.data.ul2r import DenoisingConfig, Ul2rDataset  # noqa
 from levanter.shapes import NamedShapeSpec, ShapeSpec  # noqa
 from levanter.store.cache import build_or_load_cache  # noqa
 from levanter.utils.jax_utils import key_iterator, use_cpu_device  # noqa
@@ -438,6 +439,13 @@ class SupervisedLmDatasetFormat(LmDatasetFormatBase):
     pack: bool = True
     mask_inputs: bool = True
 
+@LmDatasetFormatBase.register_subclass("ul2r")
+@dataclass(frozen=True)
+class Ul2rDatasetFormat(TextLmDatasetFormat):
+    task_configs: Dict[str, DenoisingConfig] = field(default_factory=dict)
+    task_probs: Dict[str, float] = field(default_factory=dict)
+    rng_seed: int = 37
+
 
 @dataclass(frozen=True)
 class LmDatasetSourceConfigBase(abc.ABC):
@@ -606,7 +614,7 @@ def preprocessor_for_format(
     format: LmDatasetFormatBase, tokenizer: HfTokenizer, *, enforce_eos: bool = True, enforce_bos: bool = True
 ) -> BatchProcessor[dict, dict]:
     match format:
-        case TextLmDatasetFormat(text_key=key):
+        case TextLmDatasetFormat(text_key=key) | Ul2rDatasetFormat(text_key=key):
             return BatchTokenizer(tokenizer, enforce_bos=enforce_bos, enforce_eos=enforce_eos, text_field=key)
         case ChatLmDatasetFormat(messages_field=m, single_turn=s_turn, chat_template=ct, mask_user_turns=mt):
             if s_turn:
@@ -640,6 +648,13 @@ def dataset_for_format(
                 return MultiturnChatDataset(cache, Pos, max_segments_per_example=64 if pack else 1, mask_user_turns=mask_user_turns)  # type: ignore
         case SupervisedLmDatasetFormat(pack=pack, mask_inputs=mask_inputs):
             return SupervisedDataset(cache, Pos, max_segments_per_example=64 if pack else 1, mask_inputs=mask_inputs)  # type: ignore
+        case Ul2rDatasetFormat(task_configs=task_configs, task_probs=task_probs, rng_seed=rng_seed):
+            key = jax.random.PRNGKey(rng_seed)
+            # TODO Get actual pad_token_id. Currently we only use this in ul2r_loss_mask.
+            pad_token_id = 0
+            max_segments_per_example = 64
+            slice_strategy = "left"
+            return Ul2rDataset(cache, Pos, task_configs, task_probs, key, pad_token_id, max_segments_per_example, slice_strategy)
         case _:
             raise ValueError(f"Unknown format {format}")
 

src/levanter/data/ul2r.py

@@ -892,6 +892,8 @@ def _compute_length(task_idx: jnp.ndarray, length: jnp.ndarray) -> int:
             max_segments_per_example=max_segments_per_example,
             slice_strategy=slice_strategy,
             packing_lengths=out_lengths,
+            # Reserve space for UL2R; denoising examples increase in length.
+            pad_with_zeroes=True,
         )
         self.Pos = Pos
         self.pad_token_id = pad_token_id