Language modeling on Nemotron-ClimbMix: WeLT vs. the baseline from 1.6M to 1B

benchmarks/flops.py

@@ -41,8 +41,8 @@ def step_flops(config: dict, batches: int = 20) -> dict[str, float]:
     processor = provider.processor()
     # The training examples (of a single data parallel rank), as the dataloader makes them
     texts = train_texts(provider, rank=0, world_size=1)
-    examples = (processor.process_single_example(words, seq_lengths)
-                for words, seq_lengths in pack_words(map(processor.pretokenize, texts), provider.seq_length))
+    examples = (processor.process_single_example(*packed)
+                for packed in pack_words(map(processor.pretokenize, texts), provider.seq_length))
     train_config = config.get("train") or {}
     micro_batch = train_config.get("micro_batch_size", 32)
     global_batch = train_config.get("global_batch_size", micro_batch)

tests/conftest.py

@@ -72,15 +72,17 @@ def tiny_config(tmp_path_factory):
                 "<ase>\x0eM518x529S14c20481x471S27106503x489\x0f<en> hello"]  # SignWriting
 
 
-def oracle_labels(words: list[str], seq_lengths: list[int]) -> list[str | None]:
-    """Each word's label, from the words alone: the next word of its sequence, "" after its last word, and None for
-    PAD words and for the words of a shift block before its ShiftIn (they see each other)."""
+def oracle_labels(words: list[str], seq_lengths: list[int], continues: list[bool] | None = None) -> list[str | None]:
+    """Each word's label, from the words alone: the next word of its sequence, "" after its last word (None if the
+    sequence continues in another example), and None for PAD words and for the words of a shift block before its
+    ShiftIn (they see each other)."""
     labels, start = [], 0
-    for length in seq_lengths:
+    for length, continued in zip(seq_lengths, continues or [False] * len(seq_lengths), strict=True):
         sequence, inside = words[start:start + length], False
         for i, word in enumerate(sequence):
             inside = (inside or word == SO) and word != SI
-            labels.append(None if word == PAD or inside else (sequence + [""])[i + 1])
+            last_of_continued = continued and i == length - 1
+            labels.append(None if word == PAD or inside or last_of_continued else (sequence + [""])[i + 1])
         start += length
     return labels
 

tests/test_data.py

@@ -9,12 +9,16 @@
 
 def test_pack_words_fills_blocks_and_pads():
     packed = list(pack_words([["a", "b"], ["c", "d", "e"], ["f", "g", "h", "i"]], seq_length=5))
-    assert packed == [(["a", "b", "c", "d", "e"], [2, 3]), (["f", "g", "h", "i", PAD_TOKEN], [4, 1])]
+    assert packed == [(["a", "b", "c", "d", "e"], [2, 3], [False, False]),
+                      (["f", "g", "h", "i", PAD_TOKEN], [4, 1], [False, False])]
 
 
 def test_pack_words_splits_long_sequences():
+    # Chunks overlap by one word: "d" ends the first (without a label) and starts the second
     packed = list(pack_words([list("abcdefg")], seq_length=4))
-    assert packed == [(list("abcd"), [4]), ([*"efg", PAD_TOKEN], [3, 1])]
+    assert packed == [(list("abcd"), [4], [True]), (list("defg"), [4], [False])]
+    packed = list(pack_words([list("abcde"), ["f"]], seq_length=4))
+    assert packed == [(list("abcd"), [4], [True]), ([*"de", "f", PAD_TOKEN], [2, 1, 1], [False, False, False])]
 
 
 def json_config(tmp_path, columns: dict, **kwargs) -> TextDataConfig:
@@ -35,6 +39,14 @@ def test_texts_from_a_template(tmp_path):
     assert [e["text"] for e in config.texts("validation")] == ["<s0> t0"]
 
 
+def test_texts_from_token_ids(tmp_path):
+    from transformers import AutoTokenizer
+
+    tokens = AutoTokenizer.from_pretrained("gpt2")("Hello world, שלום").input_ids
+    config = json_config(tmp_path, {"tokens": [tokens, tokens]}, dataset_tokenizer="gpt2", max_eval_samples=1)
+    assert [e["text"] for e in config.texts("validation")] == ["Hello world, שלום"]
+
+
 def test_texts_require_a_text_column(tmp_path):
     config = json_config(tmp_path, {"src": ["a", "b"]}, max_eval_samples=1)
     with pytest.raises(ValueError, match="No 'text' column"):

tests/test_metrics.py

@@ -51,11 +51,14 @@ def test_baseline_scores_the_bytes_welt_scores():
     assert tokenizer.decode(labels) == tokenizer.eos_token.join(baseline.SHIFT_BLOCK.sub("\x0e", t) for t in texts)
     baseline_bytes = sum(int(example["label_bytes"].sum()) for example in examples)
 
+    from welt_training.data_utils import pack_words
+
     processor = TextImageProcessor.create(max_word_length=32, render_images=False)
     welt_bytes = 0
-    for text in texts:
-        words = processor.pretokenize(text)
-        example = processor.process_single_example(words, [len(words)])
+    # Packed in examples of 6 words: documents are split across examples (not within a shift block), as long
+    # documents are in training
+    for packed in pack_words(map(processor.pretokenize, texts), seq_length=6):
+        example = processor.process_single_example(*packed)
         inputs = (example[k][None] for k in ("input_ids", "sequence_ids", "label_mask"))
         labels = next_word_labels(*inputs)
         welt_bytes += int(((labels != PAD_TOKEN_ID) & (labels != EOS_TOKEN_ID)).sum())

tests/test_model.py

@@ -108,7 +108,7 @@ def test_packed_attention_matches_padded_attention(megatron, causal, lengths, kv
 def packed_batch(processor, seq_length: int):
     """Two packed examples of the oracle texts (the second with PAD words), their words and sequence lengths."""
     packed = list(pack_words(map(processor.pretokenize, ORACLE_TEXTS * 20), seq_length))[:2]
-    examples = [processor.process_single_example(words, seq_lengths) for words, seq_lengths in packed]
+    examples = [processor.process_single_example(*example) for example in packed]
     return collate_fn(examples), packed
 
 
@@ -118,7 +118,7 @@ def test_latent_mask_matches_an_oracle_of_the_words(megatron, processor):
     batch, packed = packed_batch(processor, seq_length=40)
     block_mask = latent_block_mask(batch["sequence_ids"].cuda(), batch["block_ids"].cuda())
     dense = create_mask(block_mask.mask_mod, len(packed), None, 40, 40, device="cuda")[:, 0].cpu()
-    for i, (words, seq_lengths) in enumerate(packed):
+    for i, (words, seq_lengths, _) in enumerate(packed):
         torch.testing.assert_close(dense[i], oracle_mask(words, seq_lengths))
 
 
@@ -128,7 +128,7 @@ def test_masked_attention_matches_sdpa(megatron, processor):
     seq, heads, dim = 200, 4, 16  # Over one 128-word block
     batch, packed = packed_batch(processor, seq_length=seq)
     batch = {k: v.cuda() for k, v in batch.items()}
-    allowed = torch.stack([oracle_mask(words, lengths) for words, lengths in packed]).cuda()[:, None]
+    allowed = torch.stack([oracle_mask(words, lengths) for words, lengths, _ in packed]).cuda()[:, None]
     query, key, value = (torch.randn(seq, len(packed), heads, dim, device="cuda", dtype=torch.bfloat16)
                          for _ in range(3))
     mask = torch.empty(0)

tests/test_model_overfitting.py

@@ -13,8 +13,8 @@
 
 def train(model, processor, packed: bool, steps: int = 600):
     if packed:
-        batch = collate_fn([processor.process_single_example(words, lengths)
-                            for words, lengths in pack_words(map(processor.pretokenize, TRAIN_TEXTS), seq_length=7)])
+        batch = collate_fn([processor.process_single_example(*packed)
+                            for packed in pack_words(map(processor.pretokenize, TRAIN_TEXTS), seq_length=7)])
     else:
         batch = processor(TRAIN_TEXTS)
     batch = {k: v.cuda() for k, v in batch.items()}

tests/test_processor.py

@@ -109,8 +109,8 @@ def test_packed_dataset_labels_independent(processor):
         "a b",
         "c d",
     ]
-    words, seq_lengths = next(pack_words(map(processor.pretokenize, texts), seq_length=8))
-    example = processor.process_single_example(words, seq_lengths)
+    words, seq_lengths, continues = next(pack_words(map(processor.pretokenize, texts), seq_length=8))
+    example = processor.process_single_example(words, seq_lengths, continues)
     labels = [bytes(label[label > 3].tolist()).decode() if mask else None
               for label, mask in zip(labels_of(example), example["label_mask"], strict=True)]
 
@@ -194,14 +194,20 @@ def test_collate_fn_pads_every_dimension_and_keeps_dtypes():
 @pytest.mark.parametrize("seq_length", [16, 40, 128])
 def test_labels_match_an_oracle_of_the_words(text_processor, seq_length):
     """The labels the model derives (next_word_labels) on packed examples, against labels from the words alone."""
-    for words, seq_lengths in pack_words(map(text_processor.pretokenize, ORACLE_TEXTS * 2), seq_length):
-        example = text_processor.process_single_example(words, seq_lengths)
-        expected = oracle_labels(words, seq_lengths)
+    documents = [text_processor.pretokenize(text) for text in ORACLE_TEXTS * 2]
+    all_labels = []
+    for words, seq_lengths, continues in pack_words(documents, seq_length):
+        example = text_processor.process_single_example(words, seq_lengths, continues)
+        expected = oracle_labels(words, seq_lengths, continues)
+        all_labels += [label for label in expected if label is not None]
         assert example["label_mask"].tolist() == [label is not None for label in expected]
         labels = labels_of(example)[example["label_mask"]]
         tokenized = text_processor.tokenize_words([label for label in expected if label is not None]).input_ids[:, 1:]
         assert torch.equal(labels[:, :tokenized.size(1)], tokenized)
         assert (labels[:, tokenized.size(1):] == 0).all()
+    # Documents split across examples: every word's label appears exactly once, as in the unsplit documents
+    assert all_labels == [label for document in documents for label in oracle_labels(document, [len(document)])
+                          if label is not None]
 
 
 @pytest.mark.parametrize(("words", "warning", "blocks"), [

welt/processor.py

@@ -130,10 +130,11 @@ def tokenize_words(self, words: list[str]):
             truncation=True,
         )
 
-    def process_single_example(self, words: list[str], seq_lengths: list[int]):
+    def process_single_example(self, words: list[str], seq_lengths: list[int], continues: list[bool] | None = None):
         """A packed example's words (with BOS and EOS), and per word: its sequence and shift block (ids from 1, 0 for
         none), and whether it predicts the next word. PAD words (padding packed examples, each its own sequence) do not,
-        nor do the words of a shift block before its ShiftIn: they are visible to each other (bidirectional)."""
+        nor do the words of a shift block before its ShiftIn: they are visible to each other (bidirectional), nor the
+        last word of a sequence that continues in another example (see pack_words)."""
         sequence_ids = torch.repeat_interleave(torch.arange(1, len(seq_lengths) + 1), torch.tensor(seq_lengths))
         block_ids = torch.zeros(len(words), dtype=torch.long)
         label_mask = torch.tensor([word != PAD_TOKEN for word in words])
@@ -143,6 +144,8 @@ def process_single_example(self, words: list[str], seq_lengths: list[int]):
         for block, (start, end) in enumerate(blocks, start=1):
             block_ids[start:end + 1] = block
             label_mask[start:end] = False
+        ends = torch.tensor(seq_lengths).cumsum(0) - 1
+        label_mask[ends[torch.tensor(continues or [False] * len(seq_lengths), dtype=torch.bool)]] = False
         example = {"input_ids": self.tokenize_words(words).input_ids, "sequence_ids": sequence_ids,
                    "block_ids": block_ids, "label_mask": label_mask}
         if self.renderer is not None:

welt_training/baseline.py

@@ -111,7 +111,10 @@ def build(config: dict):
                                               ).to_megatron_provider(load_weights=True)
     else:
         model = transformer_provider(transformer, trust_remote_code)
-    model.vocab_size = max(model.vocab_size, len(tokenizer))  # HF embeddings may have more (padding) rows
+    # HF embeddings may have more (padding) rows. Padded to a multiple of 128 (as Pythia's 50,304) for tensor cores, and
+    # with Transformer Engine's fused cross entropy: together 2x faster steps with a vocabulary-sized LM head
+    model.vocab_size = -(-max(model.vocab_size, len(tokenizer)) // 128) * 128
+    model.cross_entropy_fusion_impl = "te"
 
     dataset = TokensDatasetProvider(tokenizer_name=model_config["tokenizer"], trust_remote_code=trust_remote_code,
                                     **config["data"])

welt_training/data.py

@@ -27,8 +27,8 @@ def build_datasets(self, context: DatasetBuildContext):
         processor = self.processor()
 
         def make_examples(texts: Iterable[str]):
-            for words, seq_lengths in pack_words(map(processor.pretokenize, texts), self.seq_length):
-                yield processor.process_single_example(words=words, seq_lengths=seq_lengths)
+            for packed in pack_words(map(processor.pretokenize, texts), self.seq_length):
+                yield processor.process_single_example(*packed)
 
         train, validation = build_iterators(self, make_examples, collate_fn,
                                             rank=parallel_state.get_data_parallel_rank(),

welt_training/data_utils.py

@@ -9,6 +9,7 @@
 import torch
 from datasets import IterableDataset, get_dataset_split_names, load_dataset
 from datasets.distributed import split_dataset_by_node
+from transformers import AutoTokenizer
 from utf8_tokenizer.tokenizer import PAD_TOKEN
 
 logger = logging.getLogger(__name__)
@@ -23,6 +24,8 @@ class TextDataConfig:
     data_files: str | list[str] | dict[str, str | list[str]] | None = None
     # A python format string over dataset columns, or [prefix, completion] which are concatenated for training
     dataset_text_template: str | list[str] | None = None
+    # Datasets of token ids (a "tokens" column, e.g. nvidia/Nemotron-ClimbMix) are decoded to texts by this tokenizer
+    dataset_tokenizer: str | None = None
     # Validation texts: the first of the validation split, or held out from the start of the train split
     max_eval_samples: int = 256
     shuffle_buffer_size: int = 10_000
@@ -48,8 +51,11 @@ def texts(self, split: str) -> IterableDataset:
         """The formatted (non-empty) texts of a split, streamed, as {"text": ...} examples."""
         template = self.dataset_text_template
         template = "".join(template) if isinstance(template, list) else template
+        tokenizer = AutoTokenizer.from_pretrained(self.dataset_tokenizer) if self.dataset_tokenizer else None
 
         def format_text(example):
+            if tokenizer is not None:
+                example = {"text": tokenizer.decode(example["tokens"])}
             if template is None and "text" not in example:
                 raise ValueError(f"No 'text' column (columns: {list(example)}), set data.dataset_text_template")
             text = example["text"] if template is None else template.format(**example)
@@ -119,20 +125,30 @@ def build_iterators(config, make_examples: Callable[[Iterable[str]], Iterator[di
     return train, cycle(validation)
 
 
-def pack_words(sequences: Iterable[list[str]], seq_length: int) -> Iterator[tuple[list[str], list[int]]]:
-    """Greedily pack word sequences, in order, into blocks of exactly seq_length words: (words, sequence lengths).
-    Longer sequences are split into seq_length chunks, each its own sequence. Each block is right-padded with PAD
-    words, each its own sequence, for constant shapes."""
-    # ponytail: a chunk's last word is trained as a document end; carry the next chunk's first word as its label
-    # (a label without an input word) if long documents matter
-    words, lengths = [], []
+def pack_words(sequences: Iterable[list[str]], seq_length: int) -> Iterator[tuple[list[str], list[int], list[bool]]]:
+    """Greedily pack word sequences, in order, into blocks of exactly seq_length words: (words, sequence lengths,
+    whether each sequence continues in a later block). Longer sequences are split into seq_length chunks, each its own
+    sequence, overlapping by one word: a chunk's last word has no label, and starts the next chunk, where it predicts
+    the next word (so every word is a label exactly once, as in the baseline's token chunks). Each block is
+    right-padded with PAD words, each its own sequence, for constant shapes."""
+    words, lengths, continues = [], [], []
+
+    def padded():
+        padding = seq_length - len(words)
+        return words + [PAD_TOKEN] * padding, lengths + [1] * padding, continues + [False] * padding
+
     for document in sequences:
-        for start in range(0, len(document), seq_length):
+        start = 0
+        while True:
             sequence = document[start:start + seq_length]
             if len(words) + len(sequence) > seq_length:
-                yield words + [PAD_TOKEN] * (seq_length - len(words)), lengths + [1] * (seq_length - len(words))
-                words, lengths = [], []
+                yield padded()
+                words, lengths, continues = [], [], []
             words += sequence
             lengths.append(len(sequence))
+            continues.append(start + seq_length < len(document))
+            if not continues[-1]:
+                break
+            start += seq_length - 1
     if words:
-        yield words + [PAD_TOKEN] * (seq_length - len(words)), lengths + [1] * (seq_length - len(words))
+        yield padded()