diff --git a/benchmarks/flops.py b/benchmarks/flops.py index 2773276..7dd24fe 100644 --- a/benchmarks/flops.py +++ b/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) diff --git a/tests/conftest.py b/tests/conftest.py index 6a44c0f..7b49344 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -72,15 +72,17 @@ def tiny_config(tmp_path_factory): "\x0eM518x529S14c20481x471S27106503x489\x0f 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 diff --git a/tests/test_data.py b/tests/test_data.py index 41f0976..27c779b 100644 --- a/tests/test_data.py +++ b/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")] == [" 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"): diff --git a/tests/test_metrics.py b/tests/test_metrics.py index 99c0539..c301c0a 100644 --- a/tests/test_metrics.py +++ b/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()) diff --git a/tests/test_model.py b/tests/test_model.py index 7f0e0f3..5c0e96a 100644 --- a/tests/test_model.py +++ b/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) diff --git a/tests/test_model_overfitting.py b/tests/test_model_overfitting.py index 485d3b5..c427c47 100644 --- a/tests/test_model_overfitting.py +++ b/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()} diff --git a/tests/test_processor.py b/tests/test_processor.py index f11b6f0..2e2cf46 100644 --- a/tests/test_processor.py +++ b/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"), [ diff --git a/welt/processor.py b/welt/processor.py index 1498268..13fada0 100644 --- a/welt/processor.py +++ b/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: diff --git a/welt_training/baseline.py b/welt_training/baseline.py index 3b63720..b34173d 100644 --- a/welt_training/baseline.py +++ b/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"]) diff --git a/welt_training/data.py b/welt_training/data.py index 29c82ff..5195ae7 100644 --- a/welt_training/data.py +++ b/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(), diff --git a/welt_training/data_utils.py b/welt_training/data_utils.py index c488e72..a8168a7 100644 --- a/welt_training/data_utils.py +++ b/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() diff --git a/welt_training/experiments/language-modeling/README.md b/welt_training/experiments/language-modeling/README.md index c0868be..5d0a89d 100644 --- a/welt_training/experiments/language-modeling/README.md +++ b/welt_training/experiments/language-modeling/README.md @@ -15,7 +15,68 @@ torchrun --nproc_per_node=1 -m welt_training.baseline welt_training/experiments/ Both log validation `bits per byte` (see the [README](../../../README.md#train)), which compare directly. -## Results +## Scaling: WeLT vs. the baseline on ClimbMix + +[nvidia/Nemotron-ClimbMix](https://huggingface.co/datasets/nvidia/Nemotron-ClimbMix) is web text stored as GPT-2 token +ids: `data.dataset_tokenizer: gpt2` decodes each streamed document back to text, which each model then segments its own +way. Five models, with latent transformers of the dimensions of Pythia's 14M, 31M, 70M, 160M, 410M and 1B +([`models/latent-*.json`](../models): 1.6M, 6.3M, 25M, 113M, 403M and 1.07B parameters without embeddings), all without +an image encoder: + +| Model | Config | Example | +|-------|--------|---------| +| Baseline: Pythia tokens, embedding matrix and LM head | [`climbmix-baseline.yaml`](climbmix-baseline.yaml) | 340 tokens | +| WeLT over Pythia tokens (the same segmentation, bytes encoder and decoder) | [`climbmix-welt-pythia.yaml`](climbmix-welt-pythia.yaml) | 340 tokens | +| WeLT over words | [`climbmix-welt.yaml`](climbmix-welt.yaml) | 256 words | +| WeLT, matched params: bytes encoder and decoder each the size of the embedding table it replaces (50,304 × hidden) | `climbmix-welt.yaml` with [`models/bytes-.json`](../models) | 256 words | +| WeLT, matched FLOPs: bytes encoder and decoder sized so that training FLOPs per step match the baseline's (within 5%) | `climbmix-welt.yaml` with [`models/bytes-flops-.json`](../models) | 256 words | + +Examples hold the same text (1.33 Pythia tokens per word). The first 32 documents are held out for validation: few +enough that each model's validation examples are a single micro batch, so all models score exactly the same bytes. +The stream is shuffled within a 10,000 document buffer, and ClimbMix stores documents by topic cluster, so long runs +drift across topics (alike for all models: comparisons at equal steps are fair). + +```shell +torchrun --nproc_per_node=1 -m welt_training.train welt_training/experiments/language-modeling/climbmix-welt.yaml \ + model.latent_transformer=welt_training/experiments/models/latent-160m.json +# Matched params: also model.bytes_encoder=welt_training/experiments/models/bytes-160m.json model.bytes_decoder=(the same) +# Matched FLOPs: bytes-flops-160m.json instead +torchrun --nproc_per_node=1 -m welt_training.baseline welt_training/experiments/language-modeling/climbmix-baseline.yaml \ + model.transformer=welt_training/experiments/models/latent-160m.json +``` + +Steps 5643 and 5651 hold outlier batches (a document of long unique strings: 2.3x the median decoder bytes) whose +activations take 2-4x the usual memory: with large byte models, lower `train.micro_batch_size` (gradient accumulation, +the same training). The 1B models train on single H100s with [`modal_1b.py`](modal_1b.py) (Modal). +[`benchmarks/flops.py`](../../../benchmarks/flops.py) counts each run's FLOPs per step. + +### Results + +Validation bits per byte after 10,000 steps (the best per latent transformer in bold): + +| Latent | Baseline | WeLT, Pythia tokens | WeLT, words | WeLT, matched params | WeLT, matched FLOPs | +|-------:|---------:|----:|----:|----:|----:| +| 1.6M | 1.238 | 1.203 | 1.194 | **1.165** | 1.242 | +| 6.3M | 1.145 | 1.146 | 1.145 | **1.097** | 1.156 | +| 25M | 1.077 | 1.088 | 1.098 | **1.043** | 1.083 | +| 113M | 1.038 | 1.057 | 1.075 | **1.020** | 1.040 | +| 403M | 1.009 | 1.024 | 1.049 | **1.007** | 1.015 | +| 1.07B | **1.024** | | | 1.038 | 1.036 | + +- Quality follows scale (FLOPs or parameters, 92% of pairs each) far more than architecture. +- At equal parameters (matched params), WeLT is ahead up to 403M, by a shrinking margin (0.072 to 0.003), and 0.014 + behind at 1B, where every model is worse than at 403M (the learning rate and text of the smaller models). +- At equal FLOPs (matched FLOPs), WeLT is slightly behind at every size (0.002-0.013): the baseline's lookups are free. +- WeLT over Pythia tokens beats WeLT over words from 25M on (0.010-0.025), also at equal FLOPs. +- Trained 30,000 steps instead of 10,000, WeLT overtakes the baseline (25M: 1.058 vs. 1.075; 113M: 1.051 vs. 1.060). +- With the 25M latent transformer, bigger byte models keep helping (1.145 with 1.1M parameters each, 1.043 with 27M, + 1.034 with 54M), and the decoder carries most of it: a 27M decoder alone gains 0.048, a 27M encoder alone 0.011. + +On a GB10, the baseline's steps are faster except with the 403M latent transformer: WeLT's small byte transformers are +bound by memory bandwidth. + +## Earlier results + 10000 steps each, on one GB10, with the BOS bytes decoder (before the latent took its place), validation bits per byte at the last step (and the best step): diff --git a/welt_training/experiments/language-modeling/climbmix-baseline.yaml b/welt_training/experiments/language-modeling/climbmix-baseline.yaml new file mode 100644 index 0000000..abd7dae --- /dev/null +++ b/welt_training/experiments/language-modeling/climbmix-baseline.yaml @@ -0,0 +1,19 @@ +# Subword causal LM baseline for climbmix-welt.yaml: the same text, transformer and training +$extends: ../machine-translation/baseline.yaml + +output_dir: ./output/lm-climbmix-baseline + +data: + dataset_name: nvidia/Nemotron-ClimbMix + dataset_config_name: null + dataset_tokenizer: gpt2 + dataset_text_template: null + seq_length: 340 # Tokens: as much text as climbmix-welt.yaml's 256 words (1.33 Pythia tokens per word) + max_eval_samples: 32 + +validation: + eval_interval: 250 + eval_iters: 1 + +logger: + wandb_project: welt-language-modeling diff --git a/welt_training/experiments/language-modeling/climbmix-welt-pythia.yaml b/welt_training/experiments/language-modeling/climbmix-welt-pythia.yaml new file mode 100644 index 0000000..0dc7829 --- /dev/null +++ b/welt_training/experiments/language-modeling/climbmix-welt-pythia.yaml @@ -0,0 +1,12 @@ +# WeLT over Pythia tokens instead of words: the baseline's segmentation (climbmix-baseline.yaml), with WeLT's bytes +# encoder and decoder instead of its embedding matrix and LM head +$extends: climbmix-welt.yaml + +output_dir: ./output/lm-climbmix-welt-pythia + +model: + pretokenizer: EleutherAI/pythia-70m + +data: + seq_length: 340 # Tokens, as climbmix-baseline.yaml + max_word_length: 32 # Pythia tokens (with BOS and EOS) diff --git a/welt_training/experiments/language-modeling/climbmix-welt.yaml b/welt_training/experiments/language-modeling/climbmix-welt.yaml new file mode 100644 index 0000000..d1006ec --- /dev/null +++ b/welt_training/experiments/language-modeling/climbmix-welt.yaml @@ -0,0 +1,25 @@ +# Language modeling on nvidia/Nemotron-ClimbMix (web text, GPT-2 token ids decoded back to text, streamed), WeLT +# without an image encoder, against climbmix-baseline.yaml +$extends: ../machine-translation/machine-translation.yaml + +output_dir: ./output/lm-climbmix-welt + +model: + image_encoder: null + +data: + dataset_name: nvidia/Nemotron-ClimbMix + dataset_config_name: null + dataset_tokenizer: gpt2 + dataset_text_template: null + seq_length: 256 + # The first documents are held out for validation: few enough that every model's validation examples are one + # micro batch, so all models score exactly the same bytes + max_eval_samples: 32 + +validation: + eval_interval: 250 + eval_iters: 1 + +logger: + wandb_project: welt-language-modeling diff --git a/welt_training/experiments/language-modeling/modal_1b.py b/welt_training/experiments/language-modeling/modal_1b.py new file mode 100644 index 0000000..379c078 --- /dev/null +++ b/welt_training/experiments/language-modeling/modal_1b.py @@ -0,0 +1,72 @@ +"""ClimbMix at Pythia-1B's dimensions on Modal H100s: the baseline, WeLT matched in parameters, WeLT matched in FLOPs +(10,000 steps each, as the local runs). Outputs and logs in the "welt-runs" volume: /out/.log, /out//. + + modal run --detach welt_training/experiments/language-modeling/modal_1b.py # all three, in parallel + modal volume get welt-runs .log . # progress (committed every 5 minutes) +""" + +import pathlib +import subprocess +import threading + +import modal + +REPO = str(pathlib.Path(__file__).resolve().parents[3]) # The repository, built with its Dockerfile +M, L = "welt_training/experiments/models", "welt_training/experiments/language-modeling" +RUNS = { + "cm-1b-baseline": f"welt_training.baseline {L}/climbmix-baseline.yaml model.transformer={M}/latent-1b.json " + "train.micro_batch_size=16", + "cm-1b-welt-matched": f"welt_training.train {L}/climbmix-welt.yaml model.latent_transformer={M}/latent-1b.json " + f"model.bytes_encoder={M}/bytes-1b.json model.bytes_decoder={M}/bytes-1b.json " + "train.micro_batch_size=8", + "cm-1b-welt-flops": f"welt_training.train {L}/climbmix-welt.yaml model.latent_transformer={M}/latent-1b.json " + f"model.bytes_encoder={M}/bytes-flops-1b.json model.bytes_decoder={M}/bytes-flops-1b.json " + "train.micro_batch_size=8", +} + +image = modal.Image.from_dockerfile(f"{REPO}/Dockerfile", context_dir=REPO).env( + { + "HF_HOME": "/hf", + "WANDB_MODE": "disabled", + "PYTHONPATH": "/app", + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + } +) +app = modal.App("welt-1b", image=image) +out = modal.Volume.from_name("welt-runs", create_if_missing=True) +hf = modal.Volume.from_name("welt-hf-cache", create_if_missing=True) + + +@app.function(gpu="H100", timeout=24 * 3600, cpu=16, memory=128 * 1024, volumes={"/out": out, "/hf": hf}) +def train(name: str) -> int: + module, *args = RUNS[name].split() + command = [ + "torchrun", + "--nproc_per_node=1", + "-m", + module, + *args, + f"output_dir=/out/{name}", + "logger.wandb_project=null", + "checkpoint.save_interval=10000", + ] + with open(f"/out/{name}.log", "w") as log: + process = subprocess.Popen(command, cwd="/app", stdout=log, stderr=subprocess.STDOUT) + stop = threading.Event() + + def commit(): # The log readable from outside while training + while not stop.wait(300): + out.commit() + + threading.Thread(target=commit, daemon=True).start() + code = process.wait() + stop.set() + out.commit() + return code + + +@app.local_entrypoint() +def main(runs: str = ",".join(RUNS)): + calls = {name: train.spawn(name) for name in runs.split(",")} + for name, call in calls.items(): + print(name, call.object_id) diff --git a/welt_training/experiments/models/bytes-14m.json b/welt_training/experiments/models/bytes-14m.json new file mode 100644 index 0000000..adda80e --- /dev/null +++ b/welt_training/experiments/models/bytes-14m.json @@ -0,0 +1,15 @@ +{ + "architectures": [ + "LlamaForCausalLM" + ], + "model_type": "llama", + "hidden_size": 256, + "intermediate_size": 640, + "num_hidden_layers": 8, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "max_position_embeddings": 2048, + "vocab_size": 256, + "rms_norm_eps": 1e-05, + "tie_word_embeddings": false +} diff --git a/welt_training/experiments/models/bytes-160m.json b/welt_training/experiments/models/bytes-160m.json new file mode 100644 index 0000000..1151bfa --- /dev/null +++ b/welt_training/experiments/models/bytes-160m.json @@ -0,0 +1,15 @@ +{ + "architectures": [ + "LlamaForCausalLM" + ], + "model_type": "llama", + "hidden_size": 640, + "intermediate_size": 1600, + "num_hidden_layers": 8, + "num_attention_heads": 10, + "num_key_value_heads": 10, + "max_position_embeddings": 2048, + "vocab_size": 256, + "rms_norm_eps": 1e-05, + "tie_word_embeddings": false +} diff --git a/welt_training/experiments/models/bytes-1b.json b/welt_training/experiments/models/bytes-1b.json new file mode 100644 index 0000000..768f3bb --- /dev/null +++ b/welt_training/experiments/models/bytes-1b.json @@ -0,0 +1,15 @@ +{ + "architectures": [ + "LlamaForCausalLM" + ], + "model_type": "llama", + "hidden_size": 1024, + "intermediate_size": 2560, + "num_hidden_layers": 8, + "num_attention_heads": 16, + "num_key_value_heads": 16, + "max_position_embeddings": 2048, + "vocab_size": 256, + "rms_norm_eps": 1e-05, + "tie_word_embeddings": false +} diff --git a/welt_training/experiments/models/bytes-31m-half.json b/welt_training/experiments/models/bytes-31m-half.json new file mode 100644 index 0000000..56bdaea --- /dev/null +++ b/welt_training/experiments/models/bytes-31m-half.json @@ -0,0 +1,15 @@ +{ + "architectures": [ + "LlamaForCausalLM" + ], + "model_type": "llama", + "hidden_size": 384, + "intermediate_size": 960, + "num_hidden_layers": 4, + "num_attention_heads": 6, + "num_key_value_heads": 6, + "max_position_embeddings": 2048, + "vocab_size": 256, + "rms_norm_eps": 1e-05, + "tie_word_embeddings": false +} diff --git a/welt_training/experiments/models/bytes-31m.json b/welt_training/experiments/models/bytes-31m.json new file mode 100644 index 0000000..b6c64ff --- /dev/null +++ b/welt_training/experiments/models/bytes-31m.json @@ -0,0 +1,15 @@ +{ + "architectures": [ + "LlamaForCausalLM" + ], + "model_type": "llama", + "hidden_size": 384, + "intermediate_size": 960, + "num_hidden_layers": 8, + "num_attention_heads": 6, + "num_key_value_heads": 6, + "max_position_embeddings": 2048, + "vocab_size": 256, + "rms_norm_eps": 1e-05, + "tie_word_embeddings": false +} diff --git a/welt_training/experiments/models/bytes-410m.json b/welt_training/experiments/models/bytes-410m.json new file mode 100644 index 0000000..4434c73 --- /dev/null +++ b/welt_training/experiments/models/bytes-410m.json @@ -0,0 +1,15 @@ +{ + "architectures": [ + "LlamaForCausalLM" + ], + "model_type": "llama", + "hidden_size": 768, + "intermediate_size": 1920, + "num_hidden_layers": 8, + "num_attention_heads": 12, + "num_key_value_heads": 12, + "max_position_embeddings": 2048, + "vocab_size": 256, + "rms_norm_eps": 1e-05, + "tie_word_embeddings": false +} diff --git a/welt_training/experiments/models/bytes-70m.json b/welt_training/experiments/models/bytes-70m.json new file mode 100644 index 0000000..136247a --- /dev/null +++ b/welt_training/experiments/models/bytes-70m.json @@ -0,0 +1,15 @@ +{ + "architectures": [ + "LlamaForCausalLM" + ], + "model_type": "llama", + "hidden_size": 512, + "intermediate_size": 1280, + "num_hidden_layers": 9, + "num_attention_heads": 8, + "num_key_value_heads": 8, + "max_position_embeddings": 2048, + "vocab_size": 256, + "rms_norm_eps": 1e-05, + "tie_word_embeddings": false +} diff --git a/welt_training/experiments/models/bytes-dec-70m.json b/welt_training/experiments/models/bytes-dec-70m.json new file mode 100644 index 0000000..e20a544 --- /dev/null +++ b/welt_training/experiments/models/bytes-dec-70m.json @@ -0,0 +1,15 @@ +{ + "architectures": [ + "LlamaForCausalLM" + ], + "model_type": "llama", + "hidden_size": 320, + "intermediate_size": 800, + "num_hidden_layers": 5, + "num_attention_heads": 5, + "num_key_value_heads": 5, + "max_position_embeddings": 2048, + "vocab_size": 256, + "rms_norm_eps": 1e-05, + "tie_word_embeddings": false +} diff --git a/welt_training/experiments/models/bytes-flops-14m.json b/welt_training/experiments/models/bytes-flops-14m.json new file mode 100644 index 0000000..c28191d --- /dev/null +++ b/welt_training/experiments/models/bytes-flops-14m.json @@ -0,0 +1,15 @@ +{ + "architectures": [ + "LlamaForCausalLM" + ], + "model_type": "llama", + "hidden_size": 128, + "intermediate_size": 320, + "num_hidden_layers": 6, + "num_attention_heads": 2, + "num_key_value_heads": 2, + "max_position_embeddings": 2048, + "vocab_size": 256, + "rms_norm_eps": 1e-05, + "tie_word_embeddings": false +} diff --git a/welt_training/experiments/models/bytes-flops-160m.json b/welt_training/experiments/models/bytes-flops-160m.json new file mode 100644 index 0000000..b136e1e --- /dev/null +++ b/welt_training/experiments/models/bytes-flops-160m.json @@ -0,0 +1,15 @@ +{ + "architectures": [ + "LlamaForCausalLM" + ], + "model_type": "llama", + "hidden_size": 384, + "intermediate_size": 960, + "num_hidden_layers": 6, + "num_attention_heads": 6, + "num_key_value_heads": 6, + "max_position_embeddings": 2048, + "vocab_size": 256, + "rms_norm_eps": 1e-05, + "tie_word_embeddings": false +} diff --git a/welt_training/experiments/models/bytes-flops-1b.json b/welt_training/experiments/models/bytes-flops-1b.json new file mode 100644 index 0000000..76bbdaa --- /dev/null +++ b/welt_training/experiments/models/bytes-flops-1b.json @@ -0,0 +1,15 @@ +{ + "architectures": [ + "LlamaForCausalLM" + ], + "model_type": "llama", + "hidden_size": 768, + "intermediate_size": 1920, + "num_hidden_layers": 9, + "num_attention_heads": 12, + "num_key_value_heads": 12, + "max_position_embeddings": 2048, + "vocab_size": 256, + "rms_norm_eps": 1e-05, + "tie_word_embeddings": false +} diff --git a/welt_training/experiments/models/bytes-flops-31m.json b/welt_training/experiments/models/bytes-flops-31m.json new file mode 100644 index 0000000..27ff323 --- /dev/null +++ b/welt_training/experiments/models/bytes-flops-31m.json @@ -0,0 +1,15 @@ +{ + "architectures": [ + "LlamaForCausalLM" + ], + "model_type": "llama", + "hidden_size": 256, + "intermediate_size": 640, + "num_hidden_layers": 3, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "max_position_embeddings": 2048, + "vocab_size": 256, + "rms_norm_eps": 1e-05, + "tie_word_embeddings": false +} diff --git a/welt_training/experiments/models/bytes-flops-410m.json b/welt_training/experiments/models/bytes-flops-410m.json new file mode 100644 index 0000000..a5bdd18 --- /dev/null +++ b/welt_training/experiments/models/bytes-flops-410m.json @@ -0,0 +1,15 @@ +{ + "architectures": [ + "LlamaForCausalLM" + ], + "model_type": "llama", + "hidden_size": 512, + "intermediate_size": 1280, + "num_hidden_layers": 8, + "num_attention_heads": 8, + "num_key_value_heads": 8, + "max_position_embeddings": 2048, + "vocab_size": 256, + "rms_norm_eps": 1e-05, + "tie_word_embeddings": false +} diff --git a/welt_training/experiments/models/bytes-flops-70m.json b/welt_training/experiments/models/bytes-flops-70m.json new file mode 100644 index 0000000..40068e8 --- /dev/null +++ b/welt_training/experiments/models/bytes-flops-70m.json @@ -0,0 +1,15 @@ +{ + "architectures": [ + "LlamaForCausalLM" + ], + "model_type": "llama", + "hidden_size": 320, + "intermediate_size": 800, + "num_hidden_layers": 4, + "num_attention_heads": 5, + "num_key_value_heads": 5, + "max_position_embeddings": 2048, + "vocab_size": 256, + "rms_norm_eps": 1e-05, + "tie_word_embeddings": false +} diff --git a/welt_training/experiments/models/latent-14m.json b/welt_training/experiments/models/latent-14m.json new file mode 100644 index 0000000..0ae8cba --- /dev/null +++ b/welt_training/experiments/models/latent-14m.json @@ -0,0 +1,15 @@ +{ + "architectures": [ + "LlamaForCausalLM" + ], + "model_type": "llama", + "hidden_size": 128, + "intermediate_size": 512, + "num_hidden_layers": 6, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "max_position_embeddings": 2048, + "vocab_size": 256, + "rms_norm_eps": 1e-05, + "tie_word_embeddings": false +} diff --git a/welt_training/experiments/models/latent-160m.json b/welt_training/experiments/models/latent-160m.json new file mode 100644 index 0000000..ac82f61 --- /dev/null +++ b/welt_training/experiments/models/latent-160m.json @@ -0,0 +1,15 @@ +{ + "architectures": [ + "LlamaForCausalLM" + ], + "model_type": "llama", + "hidden_size": 768, + "intermediate_size": 3072, + "num_hidden_layers": 12, + "num_attention_heads": 12, + "num_key_value_heads": 12, + "max_position_embeddings": 2048, + "vocab_size": 256, + "rms_norm_eps": 1e-05, + "tie_word_embeddings": false +} diff --git a/welt_training/experiments/models/latent-1b.json b/welt_training/experiments/models/latent-1b.json new file mode 100644 index 0000000..3e20a47 --- /dev/null +++ b/welt_training/experiments/models/latent-1b.json @@ -0,0 +1,15 @@ +{ + "architectures": [ + "LlamaForCausalLM" + ], + "model_type": "llama", + "hidden_size": 2048, + "intermediate_size": 8192, + "num_hidden_layers": 16, + "num_attention_heads": 8, + "num_key_value_heads": 8, + "max_position_embeddings": 2048, + "vocab_size": 256, + "rms_norm_eps": 1e-05, + "tie_word_embeddings": false +} diff --git a/welt_training/experiments/models/latent-31m.json b/welt_training/experiments/models/latent-31m.json new file mode 100644 index 0000000..43782c7 --- /dev/null +++ b/welt_training/experiments/models/latent-31m.json @@ -0,0 +1,15 @@ +{ + "architectures": [ + "LlamaForCausalLM" + ], + "model_type": "llama", + "hidden_size": 256, + "intermediate_size": 1024, + "num_hidden_layers": 6, + "num_attention_heads": 8, + "num_key_value_heads": 8, + "max_position_embeddings": 2048, + "vocab_size": 256, + "rms_norm_eps": 1e-05, + "tie_word_embeddings": false +} diff --git a/welt_training/experiments/models/latent-410m.json b/welt_training/experiments/models/latent-410m.json new file mode 100644 index 0000000..237ccba --- /dev/null +++ b/welt_training/experiments/models/latent-410m.json @@ -0,0 +1,15 @@ +{ + "architectures": [ + "LlamaForCausalLM" + ], + "model_type": "llama", + "hidden_size": 1024, + "intermediate_size": 4096, + "num_hidden_layers": 24, + "num_attention_heads": 16, + "num_key_value_heads": 16, + "max_position_embeddings": 2048, + "vocab_size": 256, + "rms_norm_eps": 1e-05, + "tie_word_embeddings": false +}