Adding configs related to DCLM

config/data/dclm_gpt_neo.yaml

@@ -0,0 +1,74 @@
+cache_dir: "gs://marin-data/tokenized/dclm/gpt_neo_tokenizer"
+tokenizer: "EleutherAI/gpt-neox-20b"
+stop_strategy: restart
+configs:
+  "dclm":
+    train_urls:
+      - gs://marin-data/datacomp/dclm-baseline-dedup-07-09/*/*/*.jsonl.zstd
+  # these are just for eval
+  "paloma/4chan":
+    validation_urls:
+      - gs://levanter-data/paloma/4chan_meta_sep/val/val*.jsonl.gz
+  "paloma/c4_100_domains":
+    validation_urls:
+      - gs://levanter-data/paloma/c4_100_domains/val/val*.jsonl.gz
+  "paloma/c4_en":
+    validation_urls:
+      - gs://levanter-data/paloma/c4_en/val/val*.jsonl.gz
+  "paloma/dolma-v1_5":
+    validation_urls:
+      - gs://levanter-data/paloma/dolma-v1_5/val/val*.jsonl.gz
+  "paloma/dolma_100_programing_languages":
+    validation_urls:
+      - gs://levanter-data/paloma/dolma_100_programing_languages/val/val*.jsonl.gz
+  "paloma/dolma_100_subreddits":
+    validation_urls:
+      - gs://levanter-data/paloma/dolma_100_subreddits/val/val*.jsonl.gz
+  "paloma/falcon-refinedweb":
+    validation_urls:
+      - gs://levanter-data/paloma/falcon-refinedweb/val/val*.jsonl.gz
+  "paloma/gab":
+    validation_urls:
+      - gs://levanter-data/paloma/gab/val/val*.jsonl.gz
+  "paloma/m2d2_s2orc_unsplit":
+    validation_urls:
+      - gs://levanter-data/paloma/m2d2_s2orc_unsplit/val/val*.jsonl.gz
+  "paloma/m2d2_wikipedia_unsplit":
+    validation_urls:
+      - gs://levanter-data/paloma/m2d2_wikipedia_unsplit/val/val*.jsonl.gz
+  "paloma/manosphere_meta_sep":
+    validation_urls:
+      - gs://levanter-data/paloma/manosphere_meta_sep/val/val*.jsonl.gz
+  "paloma/mc4":
+    validation_urls:
+      - gs://levanter-data/paloma/mc4/val/val*.jsonl.gz
+  "paloma/ptb":
+    validation_urls:
+      - gs://levanter-data/paloma/ptb/val/val*.jsonl.gz
+  "paloma/redpajama":
+    validation_urls:
+      - gs://levanter-data/paloma/redpajama/val/val*.jsonl.gz
+  "paloma/twitterAAE_HELM_fixed":
+    validation_urls:
+      - gs://levanter-data/paloma/twitterAAE_HELM_fixed/val/val*.jsonl.gz
+  "paloma/wikitext_103":
+    validation_urls:
+      - gs://levanter-data/paloma/wikitext_103/val/val*.jsonl.gz
+train_weights:
+  dclm: 1.0
+  paloma/4chan: 0.0
+  paloma/c4_100_domains: 0.0
+  paloma/c4_en: 0.0
+  paloma/dolma-v1_5: 0.0
+  paloma/dolma_100_programing_languages: 0.0
+  paloma/dolma_100_subreddits: 0.0
+  paloma/falcon-refinedweb: 0.0
+  paloma/gab: 0.0
+  paloma/m2d2_s2orc_unsplit: 0.0
+  paloma/m2d2_wikipedia_unsplit: 0.0
+  paloma/manosphere_meta_sep: 0.0
+  paloma/mc4: 0.0
+  paloma/ptb: 0.0
+  paloma/redpajama: 0.0
+  paloma/twitterAAE_HELM_fixed: 0.0
+  paloma/wikitext_103: 0.0

config/llama_1b_dclm.yaml

@@ -0,0 +1,30 @@
+data: !include data/dclm_gpt_neo.yaml
+model:  # 1B class model
+  type: llama
+  seq_len: 2048
+  hidden_dim: 2048
+  intermediate_dim: 8192
+  num_layers: 24
+  num_heads: 16
+  num_kv_heads: 16
+  use_flash_attention: True
+  flash_attention_block_size: 1024
+trainer:
+  tracker:
+    type: wandb
+    project: "marin"
+    tags: ["llama", "fineweb", "markdown"]
+
+  mp: p=f32,c=bfloat16
+  train_batch_size: 256  # 2048 * 2048 = 4,194,304
+  num_train_steps: 71526  # 300,000,000,000 / 4,194,304 = 71,526
+  steps_per_eval: 1000
+  tensor_parallel_axes: ["mlp", "heads"]
+  fsdp_axis: "embed"
+  batch_axis: "batch"
+optimizer:
+  learning_rate: 3E-3
+  weight_decay: 0.033
+  min_lr_ratio: 0.1
+  warmup: 5000
+  cooldown: 3E-5

config/llama_7b_with_dclm.yaml

@@ -0,0 +1,29 @@
+data: !include data/dclm_gpt_neo.yaml
+model:  # 7B class model
+  type: llama
+  seq_len: 2048
+  hidden_dim: 4096
+  intermediate_dim: 11008
+  num_layers: 32
+  num_heads: 32
+  num_kv_heads: 32
+  use_flash_attention: True
+  flash_attention_block_size: 1024
+trainer:
+  tracker:
+    type: wandb
+    project: "marin"
+    tags: ["dclm", "7B", "llama"]
+
+  mp: p=f32,c=bfloat16
+  train_batch_size: 2048
+  num_train_steps: 750000  # 3,000,000,000,000 / 4,000,000 = 750,000
+  steps_per_eval: 1000
+  tensor_parallel_axes: ["mlp", "heads"]
+  fsdp_axis: "embed"
+  batch_axis: "batch"
+optimizer:
+  learning_rate: 4E-4
+  weight_decay: 0.1
+  min_lr_ratio: 0.1
+  warmup: 0.01

src/levanter/main/train_lm.py

@@ -8,7 +8,7 @@
 import jax.random as jrandom
 
 import haliax as hax
-from haliax import Axis
+from haliax import Axis, Scalar
 from haliax.partitioning import named_jit, round_axis_for_partitioning
 
 import levanter
@@ -19,12 +19,29 @@
 from levanter.models.lm_model import LmConfig
 from levanter.optim import AdamConfig, OptimizerConfig
 from levanter.trainer import Trainer, TrainerConfig
+from levanter.types import ComputeLossFunction, M, X
 from levanter.utils.jax_utils import parameter_count
 
 
 logger = logging.getLogger(__name__)
 
 
+class ModuleComputeZLoss(ComputeLossFunction[M, X]):
+    """
+    Loss that just delegates to the model's compute_z_loss method.
+    """
+
+    def __call__(
+        self,
+        model,
+        *inputs: X,
+        reduction: Optional[hax.ReductionFunction] = hax.mean,
+        reduction_axis: Optional[hax.AxisSelection] = None,
+        **kwargs,
+    ) -> Scalar | hax.NamedArray:
+        return model.compute_z_loss(*inputs, reduction=reduction, reduction_axis=reduction_axis, **kwargs)
+
+
 @dataclass
 class TrainLmConfig:
     data: Union[LMDatasetConfig, LMMixtureDatasetConfig] = field(default_factory=LMDatasetConfig)
@@ -48,6 +65,7 @@ class TrainLmConfig:
 
     update_hessian_steps: int = 10
     data_seed: Optional[int] = None  # if provided, will override the data seed from the trainer
+    z_loss_weight: float = 0.0
 
 
 def main(config: TrainLmConfig):
@@ -82,11 +100,18 @@ def main(config: TrainLmConfig):
     levanter.initialize(config)
     optimizer = config.optimizer.build(config.trainer.num_train_steps)
 
+    loss_fn: Optional[ComputeLossFunction] = None
+
+    if config.z_loss_weight > 0:
+        loss_fn = ModuleComputeZLoss()
+    else:
+        loss_fn = None  # It will be automatically set to the default loss function in the model
+
     # Using the trainer as a context manager does 3 things:
     # 1. Sets the device mesh
     # 2. Sets the axis mapping (for fsdp)
     # 3. Sets the global metrics tracker
-    with Trainer(config.trainer, optimizer) as trainer:
+    with Trainer(config.trainer, optimizer, loss_fn) as trainer:
         # randomness in jax is tightly controlled by "keys" which are the states of the random number generators
         # this makes deterministic training pretty easy
         seed = config.trainer.seed

src/levanter/models/lm_model.py

@@ -9,8 +9,10 @@
 import haliax as hax
 from haliax import Axis, NamedArray
 from haliax.nn import cross_entropy_loss
+from haliax.nn.loss import maybe_reduce_loss
 
 from levanter.models.attention import AttentionMask
+from levanter.models.loss import cross_entropy_and_logsumexp_penalty
 
 
 LmConfigT = TypeVar("LmConfigT", bound="LmConfig")
@@ -137,6 +139,32 @@ def compute_loss(
 
         return loss
 
+    def compute_z_loss(
+        self,
+        example: LmExample,
+        z_loss_weight,
+        *,
+        key=None,
+        reduction: Optional[hax.ReductionFunction] = hax.mean,
+        reduction_axis: Optional[hax.AxisSelection] = None,
+    ) -> jnp.ndarray | NamedArray:
+        """
+        Computes the cross-entropy loss for a language modeling example with z_loss.
+        If reduction is not None, the loss is reduced
+        across the reduction axis (with reduction_axis=None meaning all axes). If reduction is None, the loss is not
+        reduced, and the result is a named array with axes (*batch axes, sequence_length).
+        """
+        logits = self(example.tokens, example.attn_mask, key=key)
+        # TODO: would be nice if we made the dtype configurable
+        logits = logits.astype(jnp.float32)
+        targets = hax.roll(example.tokens, -1, axis=self.Pos.name)
+        target_y = hax.nn.one_hot(targets, self.Vocab, dtype=logits.dtype)
+        loss = cross_entropy_and_logsumexp_penalty(
+            logits, self.Vocab, target_y, logsumexp_weight=self.config.z_loss_weight
+        )
+        loss = maybe_reduce_loss(loss, reduction=reduction, reduction_axis=reduction_axis, where=example.loss_mask)
+        return loss
+
     @property
     def vocab_size(self) -> int:
         return self.Vocab.size