off by one

stanford-crfm · Nov 15, 2024 · c7c5f70 · c7c5f70
1 parent 3040956
commit c7c5f70
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/config/gpt2_nano_harness.yaml b/config/gpt2_nano_harness.yaml
@@ -19,7 +19,7 @@ trainer:
     save_interval: 5m
 
   per_device_parallelism: -1
-  train_batch_size: 32
+  train_batch_size: 4
 
   tensor_parallel_axes: ["mlp", "heads"]
   fsdp_axis: "embed"

diff --git a/src/levanter/models/lm_model.py b/src/levanter/models/lm_model.py
@@ -57,7 +57,7 @@ def from_prompt_and_completion(
         all_causal: bool = True,
     ) -> "LmExample":
         # mask out the prompt tokens
-        loss_mask = hax.arange(Pos) >= prompt_length
+        loss_mask = hax.arange(Pos) >= prompt_length - 1
         # also mask out the last token
         loss_mask *= 1 - hax.nn.one_hot(-1, Pos, dtype=jnp.float32)