diff --git a/README.md b/README.md index ea448c8..14e2cc3 100644 --- a/README.md +++ b/README.md @@ -7,8 +7,9 @@ We present a language model that replaces conventional subword tokenization with a dual representation of text as both bytes and rendered images of "words", allowing the model to directly process the visual and symbolic -structure of language. It segments text into whitespace-delimited units, encodes each as images and byte sequences, -and passes them through an encoder–decoder pipeline: +structure of language. It segments text into words (with [`sign/words-segmentation`](https://github.com/sign/words-segmentation) +by default, or any HuggingFace tokenizer), encodes each as a byte sequence and/or an image, and passes them through an +encoder–decoder pipeline: a bytes encoder, an image encoder, a large latent transformer, and a bytes decoder. At inference, predicted bytes are rendered back into images (with an image encoder), closing the loop for the next prediction. This design could make learning and inference cheaper and faster for non-English languages, @@ -25,9 +26,11 @@ has an embedding, and the latent transformer sees one vector per word. - **Bytes Encoder** - A transformer over the bytes of each word, used bidirectionally. Its first (BOS) output is the word embedding. - **Image Encoder** - A transformer over the 16x16 patches of each rendered word - ([NaViT](https://arxiv.org/abs/2307.06304)-style, variable width, used bidirectionally). Its CLS output is the word embedding. + ([NaViT](https://arxiv.org/abs/2307.06304)-style: native size, with learned row and column embeddings, packed + without padding; used bidirectionally). Its CLS output is the word embedding. - **Latent Transformer** - A causal transformer over the word embeddings (recommended: large). -- **Bytes Decoder** - A causal transformer generating each next word's bytes from its latent vector (recommended: small). +- **Bytes Decoder** - A causal transformer generating each next word's bytes from its latent vector, which takes the + place of BOS: its input is the latent, then the word's bytes, and it predicts the bytes and EOS (recommended: small). The bytes encoder, latent transformer, bytes decoder and patch image encoder are each a causal LM architecture [supported by Megatron-Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge#supported-models) @@ -125,7 +128,8 @@ and reports exact match, chrF and generated words per second. ## Baseline `welt_training.baseline` trains a causal LM over a standard tokenizer on the same data, reporting comparable bits per -byte; see [machine translation](welt_training/experiments/machine-translation). +byte; see [machine translation](welt_training/experiments/machine-translation) and +[language modeling](welt_training/experiments/language-modeling). ## Training Speed diff --git a/assets/architecture.png b/assets/architecture.png index fa32ce9..b6bbdb9 100644 Binary files a/assets/architecture.png and b/assets/architecture.png differ diff --git a/benchmarks/README.md b/benchmarks/README.md index 1ad7576..e855dd8 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -36,6 +36,8 @@ harness, and the time per step of every optimization attempt, are in git history transformer's arbitrary (packed + bidirectional shift block) mask. 4. **Encode unique words** (1.3x on MT): the word encoders' output only depends on the word, and only ~31% of the words in a batch are distinct, so each distinct word is encoded once and its embedding is shared (exact). +5. **Packed patches** (SignWriting 404 -> 262 ms): the image encoder packs each batch's patches, with patch row and + column embeddings, instead of padding every sign (28 patches on average) to the batch's largest (~105). Not kept: padded SDPA / matmul attention for the short sequences (memory-bound on GB10), FP8 (MXFP8 is unsupported on SM12x; tensorwise FP8 needs token counts padded to multiples of 8, and matmuls are a minority of the step time here). @@ -43,7 +45,8 @@ SM12x; tensorwise FP8 needs token counts padded to multiples of 8, and matmuls a ### Utilization [`flops.py`](flops.py) counts the FLOPs of each transformer on its real (packed) tokens: "model" FLOPs encode every word -occurrence, "hardware" FLOPs encode each distinct word once. +occurrence, "hardware" FLOPs encode each distinct word once +(`python benchmarks/flops.py --ms_per_step `). | Config | ms / step | Model TFLOPs / step | Model TFLOP/s | Hardware TFLOP/s | |--------|----------:|-------------------:|--------------:|-----------------:| @@ -112,17 +115,18 @@ validation prefixes), recording a row of [`tasks.csv`](tasks.csv). | Task | Steps | ms / step | Train time | Val. bits / byte | Val. word acc. | Gen. exact match | Gen. chrF | Gen. words / s | |------|------:|----------:|-----------:|-----------------:|---------------:|-----------------:|----------:|---------------:| -| **string-repetition**: repeat an English sentence, pretrained tiny LMs | 1500 | 66 | 1.6 min | 0.014 | 99.3% | 93.8% | 98.7 | 856 | -| **ocr**: write a sentence seen only as rendered word images | 3000 | 74 | 3.7 min | 0.014 | 99.1% | 88.3% | 97.1 | 1055 | -| **letter-count**: count the letters of a word | 3000 | 64 | 3.2 min | 0.0001 | 100% | 100% | 100 | 1378 | -| **machine-translation**: English to Hebrew, from scratch, image + bytes encoders | 10000 | 198 | 33 min | 0.547 | 67.7% | 8.2% | 45.4 | 536 | -| **machine-translation**, bytes encoder only | 10000 | 170 | 28 min | 0.537 | 68.0% | 9.0% | 46.5 | 616 | -| **signed-to-spoken**: SignWriting to text ([signbank-plus](https://huggingface.co/datasets/sign/signbank-plus)), from scratch, bytes encoder | 10000 | 224 | 37 min | 1.253 | 63.1% | 2.3% | 11.1 | 378 | -| **signed-to-spoken**, bytes + image encoders (rendered SignWriting) | 10000 | 264 | 44 min | 1.270 | 63.7% | 4.3% | 14.2 | 304 | -| **signed-to-spoken**, image encoder only | 10000 | 166 | 28 min | 1.303 | 62.9% | 0.8% | 10.5 | 267 | - -All tasks train with Muon (their configs), on main after the migration. Generation is greedy, on 256 validation -examples, with vLLM (batched over all examples). Reproduce a row with e.g. +| **string-repetition**: repeat an English sentence, pretrained tiny LMs | 1500 | 59 | 1.5 min | 0.008 | 99.3% | 93.8% | 99.2 | 548 | +| **ocr**: write a sentence seen only as rendered word images | 3000 | 62 | 3.1 min | 0.009 | 98.8% | 87.5% | 97.8 | 449 | +| **letter-count**: count the letters of a word | 3000 | 65 | 3.2 min | 0.0001 | 100.0% | 100.0% | 100.0 | 1373 | +| **machine-translation**: English to Hebrew, from scratch, image + bytes encoders | 10000 | 182 | 30.3 min | 0.531 | 68.5% | 9.8% | 47.6 | 910 | +| **machine-translation**, bytes encoder only | 10000 | 160 | 26.8 min | 0.528 | 68.0% | 6.2% | 46.2 | 867 | +| **signed-to-spoken**: SignWriting to text ([signbank-plus](https://huggingface.co/datasets/sign/signbank-plus)), from scratch, bytes encoder | 10000 | 209 | 34.9 min | 1.244 | 63.1% | 2.7% | 10.4 | 560 | +| **signed-to-spoken**, bytes + image encoders (rendered SignWriting) | 10000 | 248 | 41.4 min | 1.246 | 63.4% | 2.3% | 12.2 | 638 | +| **signed-to-spoken**, image encoder only | 10000 | 150 | 24.9 min | 1.277 | 63.1% | 2.3% | 10.8 | 790 | + +All tasks train with Muon (their configs). Steps are `train.train_iters` overrides of the configs' 10000 +(string-repetition: 1500; ocr, letter-count: 3000). Generation is greedy, on 256 validation +examples (the easy tasks' configs hold out 32), with vLLM (batched over all examples). Reproduce a row with e.g. `WANDB_MODE=disabled benchmarks/run_task.sh welt_training/experiments/easy-tasks/string-repetition.yaml output/string-repetition train.train_iters=1500`. @@ -136,23 +140,17 @@ dataset is small (57k packed examples): every configuration's validation loss is then overfits. Dropout in the latent transformer slows this: `hidden_dropout: 0.2` (the config's) gives the best last step (1.270 bits per byte with both encoders, vs. 1.303 with 0.1 and 1.276 with 0.3; weight decay 0.1 did not help). -The image encoder (over SignWriting rendered by pixel-renderer, from scratch) helps the best checkpoint: both encoders -reach 1.212 bits per byte (step 7000) vs. 1.239 with bytes only and 1.263 with the image encoder only. By the last step -all have overfitted (bytes only is lowest there, 1.253 vs. 1.270), while both encoders generate best (4.3% exact -matches, chrF 14.2); but generation, on 256 examples, is too noisy to rank these (chrF moves by up to 5 between checkpoints). -Modality dropout (training with one encoder's embeddings dropped at random) was removed: with zeros and rescaling, or -a learned embedding standing in for the dropped encoder, it was worse than none (1.25 vs. 1.23 bits per byte, and -under half the exact matches). - -An earlier comparison (bytes only 0.953 bits per byte, better than with images) was wrong: configs without images read -signs split into fragments by words-segmentation 0.0.5 from the datasets cache, which did not invalidate on the -upgrade to 0.0.6 (data is now streamed and split into words on the fly, without a cache). +The image encoder (over SignWriting rendered by pixel-renderer, from scratch) barely helps: at their best checkpoints, +both encoders reach 1.174 bits per byte (step 7000) vs. 1.178 with bytes only (step 6000) and 1.258 with the image +encoder only. By the last step all have overfitted (bytes only is lowest there, 1.244 vs. 1.246); both encoders +generate with the best chrF (12.2), bytes only with the most exact matches (2.7%), but generation, on 256 examples, is +too noisy to rank these (chrF moves by up to 5 between checkpoints). +There is no modality dropout (one encoder's embeddings dropped at random): it was worse than none (1.25 vs. 1.23 bits +per byte, and under half the exact matches). Images do not help machine translation: without the image encoder (`model.image_encoder=null`), English to Hebrew -reaches 0.537 bits per byte and chrF 46.5 (vs. 0.547 and 45.4 with it), 14% faster -per step. The image encoder packs each batch's patches (no padding to the largest image) and embeds patch rows and -columns: before, padding signs (28 patches on average) to the batch's largest (~105) copied 425 MB per batch to the -GPU and made a SignWriting step 404 ms; packed, 41 MB and 262 ms. +reaches 0.528 bits per byte and chrF 46.2 (vs. 0.531 and 47.6 with it, within the noise of generation), 12% faster +per step. Each task trains its own models and data (see its config): string-repetition, ocr and letter-count use tiny transformers and short sequences, so their steps are faster than the bench config's 127 ms, while machine diff --git a/benchmarks/tasks.csv b/benchmarks/tasks.csv index d36b9eb..20d5804 100644 --- a/benchmarks/tasks.csv +++ b/benchmarks/tasks.csv @@ -1,9 +1,9 @@ task,train_iters,ms_per_step,val_loss,val_bits_per_byte,val_byte_accuracy,val_word_accuracy,gen_exact_match,gen_chrf,gen_words_per_second -string-repetition,1500,65.6,0.0083,0.0144,0.9987,0.9926,0.9375,98.68,855.9 -ocr,3000,73.6,0.0081,0.0140,0.9979,0.9905,0.8828,97.09,1055.1 -letter-count,3000,64.0,0.0000,0.0001,1.0000,1.0000,1.0000,100.00,1377.9 -machine-translation,10000,198.3,0.3273,0.5469,0.9082,0.6767,0.0820,45.37,536.4 -machine-translation (bytes),10000,169.6,0.3215,0.5369,0.9086,0.6799,0.0898,46.52,616.2 -signed-spoken (bytes),10000,224.0,0.6866,1.2526,0.8091,0.6313,0.0234,11.14,377.8 -signed-spoken (bytes + image),10000,264.1,0.6961,1.2698,0.8133,0.6371,0.0430,14.20,304.0 -signed-spoken (image),10000,166.0,0.7193,1.3025,0.8062,0.6289,0.0078,10.50,266.8 +string-repetition,1500,59.3,0.0045,0.0078,0.9987,0.9927,0.9375,99.23,548.3 +ocr,3000,62.1,0.0051,0.0088,0.9979,0.9879,0.8750,97.82,449.2 +letter-count,3000,64.7,0.0001,0.0001,1.0000,1.0000,1.0000,100.00,1373.1 +machine-translation,10000,181.8,0.3180,0.5312,0.9113,0.6854,0.0977,47.62,909.8 +machine-translation (bytes),10000,160.5,0.3159,0.5283,0.9101,0.6795,0.0625,46.23,867.0 +signed-spoken (bytes),10000,209.4,0.6817,1.2435,0.8162,0.6306,0.0273,10.44,559.8 +signed-spoken (bytes + image),10000,248.1,0.6830,1.2458,0.8206,0.6340,0.0234,12.17,638.0 +signed-spoken (image),10000,149.6,0.7046,1.2773,0.8044,0.6313,0.0234,10.78,790.4 diff --git a/welt_training/README.md b/welt_training/README.md index 162919c..0660647 100644 --- a/welt_training/README.md +++ b/welt_training/README.md @@ -72,13 +72,16 @@ from the start of the train split. `dataset_text_template` is a Python format string over the dataset's columns; without it, the `text` column is used (a missing `text` column raises). As `[prefix, completion]`, both parts are concatenated for training, and -`welt-evaluate` generates the completion from the prefix. +`welt-evaluate` generates the completion from the prefix. Datasets of token ids (a `tokens` column, e.g. +`nvidia/Nemotron-ClimbMix`) are decoded back to text with `dataset_tokenizer` (e.g. `gpt2`). Training streams endlessly (epochs reshuffle, with a `shuffle_buffer_size` buffer and `seed`), split across data parallel ranks, and in each rank across the dataloader workers (`num_workers`), which make the examples on the fly: texts are prefixed with BOS, split into words, and packed in order into examples of exactly `seq_length` words. -Longer texts are split into `seq_length`-word chunks, each its own sequence, and the rest is padded. -With the default pretokenizer, words longer than `max_word_length - 2` bytes are split. +Longer texts are split into `seq_length`-word chunks, each its own sequence, overlapping by one word (a chunk's last +word predicts nothing, and predicts the next word where it starts the next chunk), and the rest is padded. +With the default pretokenizer, words longer than `max_word_length - 2` bytes are split; with other pretokenizers, +they are truncated. Within the model, the bytes of all words (and the patches of rendered words) are packed without padding, so the encoders and the bytes decoder only compute on real bytes. Resuming from a checkpoint restarts the stream. @@ -102,9 +105,6 @@ Megatron's optimizers, selected by `optimizer.optimizer`: `muon` (default), `ada Adam (see [benchmarks](../benchmarks/README.md#tasks)). Its layer-wise distributed variant (`dist_muon`) fails. Their hyperparameters are `OptimizerConfig` fields (e.g. `optimizer.muon_momentum`). -Only UTF-8 bytes are supported; the previous HuggingFace Trainer implementation is at the `huggingface-transformers` -git tag. - ## Experimenting - **A new task**: a YAML with `$extends: ../easy-tasks/string-repetition.yaml` (the path is relative to the new file), diff --git a/welt_training/evaluate.py b/welt_training/evaluate.py index 4551912..c0028c9 100644 --- a/welt_training/evaluate.py +++ b/welt_training/evaluate.py @@ -26,7 +26,8 @@ def evaluate(config_path: str, url: str, max_samples: int = 256, max_generated_w prefixes = [prefix_template.format(**example) for example in validation] references = [completion_template.format(**example) for example in validation] - generate(url, prefixes, max_generated_words=2) # Warmup, with the batch size of the measurement + # Warmup, as the measurement: shorter generations leave decode kernels (e.g. Triton's) to compile while measured + generate(url, prefixes, max_generated_words=max_generated_words) start = time.perf_counter() result = generate(url, prefixes, max_generated_words=max_generated_words) elapsed = time.perf_counter() - start diff --git a/welt_training/experiments/language-modeling/README.md b/welt_training/experiments/language-modeling/README.md index 5d0a89d..b676bb6 100644 --- a/welt_training/experiments/language-modeling/README.md +++ b/welt_training/experiments/language-modeling/README.md @@ -78,7 +78,9 @@ 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): +10000 steps each, on one GB10, with the BOS bytes decoder (before the latent took its place) and the baseline's +unpadded vocabulary with Megatron's cross entropy (padded and fused, its steps now take half the time), validation bits +per byte at the last step (and the best step): | Data | WeLT | Baseline | ms / step (WeLT, baseline) | |------|-----:|---------:|---------------------------:| @@ -86,5 +88,4 @@ bound by memory bandwidth. | wikitext-103 | 1.199 | **1.165** | **197**, 226 | The baseline's sequence lengths match WeLT's text per example: Pythia's tokenizer splits opus-100 (half Hebrew) into -2.7 tokens per word, but wikitext into 1.1. WeLT wins where subword tokens are short (2.1x faster per step, and -better), and the baseline where they are nearly whole words. +2.7 tokens per word, but wikitext into 1.1. WeLT wins where subword tokens are short (better, and then 2.1x faster per step), and the baseline where they are nearly whole words. diff --git a/welt_training/experiments/pretrain/README.md b/welt_training/experiments/pretrain/README.md new file mode 100644 index 0000000..65a1599 --- /dev/null +++ b/welt_training/experiments/pretrain/README.md @@ -0,0 +1,11 @@ +# Pretraining + +[`pile-pretrain-70m-no-image.yaml`](pile-pretrain-70m-no-image.yaml) pretrains a 70m latent transformer, bytes encoder +only, on the Pile (`monology/pile-uncopyrighted`, streamed), for Pythia's 143k steps. [`test-run.yaml`](test-run.yaml) +runs it for 100 small steps, as a check of the pipeline. + +```shell +torchrun --nproc_per_node=1 -m welt_training.train welt_training/experiments/pretrain/test-run.yaml +``` + +For WeLT against a subword baseline at several sizes, see [language modeling](../language-modeling).