From 9515114437278e10d3d838bac7627fa70b135f8f Mon Sep 17 00:00:00 2001 From: autoresearch Date: Tue, 13 Oct 2026 03:52:31 -0400 Subject: [PATCH] Audit: evaluate every validation example once; shift blocks cut by packing; reloaded processors split words as in training; data and serving fixes; tests Validation: - Every validation example is scored exactly once (copies with score 0 fill an evaluation's remaining micro batches); an evaluation that cannot cover them fails instead of seeing only the first. Before, models saw only the examples that fit in validation.eval_iters: on ClimbMix, WeLT the first 64 of 80 and the baseline the first 64 of 72, different texts. ClimbMix configs: eval_iters 2. - Data split across GPUs before skipping the held out validation examples (they leaked into training on multiple GPUs with sharded datasets); dataloader workers format (decode) only their share of the stream. Correctness: - Shift blocks cut by pack_words: a sequence ending (beginning) within a block keeps it to its last (from its first) word, so WeLT scores the bytes the baseline scores; a ShiftOut within a block belongs to it, as the baseline's regex. - Processors loaded with from_pretrained split words as in training (words-segmentation does not save max_bytes), and keep trust_remote_code; other pretokenizers' long tokens are split rather than truncated. - Baseline: model options set the transformer's (unknown ones raise); data.trust_remote_code no longer clashes. - Generation: each text samples with its own seed; generated words are counted, not re-segmented; only a context too long (ContextTooLongError) answers 400; welt-generate --seed. - flops.py reuses the provider's make_examples and build_config's batch defaults. Tests: oracle labels across split documents (also within shift blocks), byte parity at short sequence lengths (in CI, with the token chunking moved to data_utils), dataloader workers' split, NUL and empty texts, the generation loop, configs' model files, reloaded processors; the busy-server test no longer hangs. CI: Docker build on pyproject changes, ruff pinned. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015Xm62qoeaTunFSU5GYCQ1f --- .github/workflows/docker-build.yaml | 1 + .github/workflows/lint.yaml | 4 +- .github/workflows/publish-docker.yaml | 2 +- benchmarks/flops.py | 11 +-- tests/conftest.py | 22 +++-- tests/test_config.py | 32 ++++++- tests/test_data.py | 87 +++++++++++++++-- tests/test_inference.py | 44 +++++++++ tests/test_metrics.py | 40 -------- tests/test_processor.py | 47 +++++---- tests/test_server.py | 17 ++-- tests/test_train.py | 3 +- welt/__init__.py | 2 + welt/inference.py | 16 +++- welt/processor.py | 57 +++++++---- welt/server.py | 13 ++- welt_training/README.md | 7 +- welt_training/baseline.py | 64 ++++--------- welt_training/data.py | 16 ++-- welt_training/data_utils.py | 96 ++++++++++++++----- .../experiments/language-modeling/README.md | 8 +- .../language-modeling/climbmix-baseline.yaml | 2 +- .../language-modeling/climbmix-welt.yaml | 6 +- welt_training/train.py | 10 +- 24 files changed, 395 insertions(+), 212 deletions(-) create mode 100644 tests/test_inference.py diff --git a/.github/workflows/docker-build.yaml b/.github/workflows/docker-build.yaml index 8670480..23943e0 100644 --- a/.github/workflows/docker-build.yaml +++ b/.github/workflows/docker-build.yaml @@ -9,6 +9,7 @@ on: pull_request: paths: - 'Dockerfile' + - 'pyproject.toml' jobs: docker-build: diff --git a/.github/workflows/lint.yaml b/.github/workflows/lint.yaml index 2c24a9f..f8f6b62 100644 --- a/.github/workflows/lint.yaml +++ b/.github/workflows/lint.yaml @@ -16,4 +16,6 @@ jobs: - uses: actions/checkout@v5 - name: Lint code - uses: astral-sh/ruff-action@v3 \ No newline at end of file + uses: astral-sh/ruff-action@v3 + with: + version: "0.15.8" # Pinned: a new ruff release adds rules diff --git a/.github/workflows/publish-docker.yaml b/.github/workflows/publish-docker.yaml index f6b49a4..9bae373 100644 --- a/.github/workflows/publish-docker.yaml +++ b/.github/workflows/publish-docker.yaml @@ -24,7 +24,7 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@v4 + uses: actions/checkout@v5 # Reclaim ~50โ€“80 GB on the GitHub runner - name: Free disk space diff --git a/benchmarks/flops.py b/benchmarks/flops.py index 7dd24fe..63806c5 100644 --- a/benchmarks/flops.py +++ b/benchmarks/flops.py @@ -14,7 +14,7 @@ from utf8_tokenizer.tokenizer import PAD_TOKEN_ID from welt.processor import collate_fn, next_word_labels -from welt_training.data_utils import pack_words, train_texts +from welt_training.data_utils import train_texts from welt_training.extendable_yaml import load_yaml from welt_training.train import build_dataset_provider @@ -38,14 +38,11 @@ def transformer_flops(config, lengths) -> float: def step_flops(config: dict, batches: int = 20) -> dict[str, float]: model, data = config["model"], config["data"] provider = build_dataset_provider(model, data) - 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(*packed) - for packed in pack_words(map(processor.pretokenize, texts), provider.seq_length)) + examples = provider.make_examples(train_texts(provider, rank=0, world_size=1)) 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) + micro_batch = train_config.get("micro_batch_size", 32) # build_config's defaults + global_batch = train_config.get("global_batch_size", 32) configs = {name: AutoConfig.from_pretrained(model[name], trust_remote_code=model.get("trust_remote_code", False)) for name in ["bytes_encoder", "image_encoder", "latent_transformer", "bytes_decoder"] if model.get(name)} diff --git a/tests/conftest.py b/tests/conftest.py index 7b49344..1fc3788 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -75,10 +75,12 @@ def tiny_config(tmp_path_factory): 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).""" + ShiftIn (they see each other). A sequence cut from a longer document can begin within a block (a ShiftIn before + any ShiftOut).""" labels, start = [], 0 for length, continued in zip(seq_lengths, continues or [False] * len(seq_lengths), strict=True): - sequence, inside = words[start:start + length], False + sequence = words[start:start + length] + inside = SI in sequence and SO not in sequence[:sequence.index(SI)] for i, word in enumerate(sequence): inside = (inside or word == SO) and word != SI last_of_continued = continued and i == length - 1 @@ -89,17 +91,21 @@ def oracle_labels(words: list[str], seq_lengths: list[int], continues: list[bool def oracle_mask(words: list[str], seq_lengths: list[int]) -> torch.Tensor: """(L, L) the words each word sees in the latent transformer, from the words alone: itself and the earlier words - of its sequence, and every word of its shift block (ShiftOut to ShiftIn).""" + of its sequence, and every word of its shift block (ShiftOut to ShiftIn; a sequence cut from a longer document + can begin or end within one).""" allowed = torch.zeros(len(words), len(words), dtype=torch.bool) start = 0 for length in seq_lengths: - shift_out = None - for i in range(start, start + length): + shift_out, end = None, start + length + for i in range(start, end): allowed[i, start:i + 1] = True - if words[i] == SO: + if words[i] == SO and shift_out is None: shift_out = i - elif words[i] == SI and shift_out is not None: - allowed[shift_out:i + 1, shift_out:i + 1] = True + elif words[i] == SI: + block_start = start if shift_out is None else shift_out + allowed[block_start:i + 1, block_start:i + 1] = True shift_out = None + if shift_out is not None: + allowed[shift_out:end, shift_out:end] = True start += length return allowed diff --git a/tests/test_config.py b/tests/test_config.py index 7de6b5d..2230459 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,9 +1,12 @@ import glob +from pathlib import Path import pytest from welt_training.extendable_yaml import load_yaml +ROOT = Path(__file__).resolve().parents[1] + def test_nested_extends_and_overrides(tmp_path): (tmp_path / "base.yaml").write_text("model:\n a: 1\n b: 2\ndata:\n seq_length: 128\n") @@ -15,23 +18,34 @@ def test_nested_extends_and_overrides(tmp_path): "output_dir": "out", "train": {"train_iters": 5}, "optimizer": {"lr": 1e-4}} -CONFIGS = sorted(glob.glob("welt_training/experiments/*/*.yaml") + ["benchmarks/welt-bench.yaml"]) +CONFIGS = sorted([*glob.glob(str(ROOT / "welt_training/experiments/*/*.yaml")), + str(ROOT / "benchmarks/welt-bench.yaml")]) @pytest.mark.parametrize("path", CONFIGS) -def test_experiment_configs_build(path): +def test_experiment_configs_reference_existing_models(path): + """Model configs given as JSON files exist (paths are relative to the repository).""" + for value in load_yaml(path)["model"].values(): + if isinstance(value, str) and value.endswith(".json"): + assert (ROOT / value).exists(), value + + +@pytest.mark.parametrize("path", CONFIGS) +def test_experiment_configs_build(path, monkeypatch): """Every shipped config builds a Megatron-Bridge config (unknown options raise).""" pytest.importorskip("megatron.bridge", reason="Requires the NeMo container") from welt_training import baseline, train config = load_yaml(path) + monkeypatch.chdir(ROOT) # Model JSON paths are relative to the repository (baseline if "transformer" in config["model"] else train).build(config) # baseline: a causal LM -def test_null_sections_and_model_overrides(): +def test_null_sections_and_model_overrides(monkeypatch): pytest.importorskip("megatron.bridge", reason="Requires the NeMo container") from welt_training import train + monkeypatch.chdir(ROOT) config = load_yaml("welt_training/experiments/easy-tasks/string-repetition.yaml", ["optimizer=null", "model.hidden_dropout=0.2"]) cfg, _ = train.build(config) @@ -40,3 +54,15 @@ def test_null_sections_and_model_overrides(): train.build(load_yaml("welt_training/experiments/easy-tasks/string-repetition.yaml", ["train.typo=1"])) with pytest.raises(ValueError, match="Unknown config sections"): train.build(load_yaml("welt_training/experiments/easy-tasks/string-repetition.yaml", ["optimiser.lr=0.1"])) + + +def test_baseline_model_overrides(monkeypatch): + pytest.importorskip("megatron.bridge", reason="Requires the NeMo container") + from welt_training import baseline + + monkeypatch.chdir(ROOT) + path = "welt_training/experiments/language-modeling/climbmix-baseline.yaml" + cfg, _ = baseline.build(load_yaml(path, ["model.hidden_dropout=0.2"])) + assert cfg.model.hidden_dropout == 0.2 + with pytest.raises(ValueError, match="Unknown model option"): + baseline.build(load_yaml(path, ["model.typo=1"])) diff --git a/tests/test_data.py b/tests/test_data.py index 27c779b..01aba3a 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -1,9 +1,11 @@ from itertools import islice import pytest +import torch from datasets import Dataset -from utf8_tokenizer.tokenizer import PAD_TOKEN +from utf8_tokenizer.tokenizer import EOS_TOKEN_ID, PAD_TOKEN, PAD_TOKEN_ID +from welt_training import data_utils from welt_training.data_utils import TextDataConfig, build_iterators, pack_words, train_texts @@ -68,11 +70,82 @@ def test_build_iterators_batches_examples(tmp_path): def make_examples(texts): yield from ({"text": text} for text in texts) - train, validation = build_iterators(config, make_examples, list, rank=0, world_size=1) + def collate(examples): + return {"texts": [example["text"] for example in examples]} + + train, validation = build_iterators(config, make_examples, collate, rank=0, world_size=1) batches = list(islice(train, 4)) - assert all(len(batch) == 2 for batch in batches) - assert {e["text"] for batch in batches for e in batch} <= {f"t{i}" for i in range(3, 10)} - # Each evaluation sees the validation examples (repeated to fill its micro batches), the same ones every time + assert all(len(batch["texts"]) == 2 for batch in batches) + assert {text for batch in batches for text in batch["texts"]} <= {f"t{i}" for i in range(3, 10)} + # Each evaluation sees every validation example once, filled up with copies that count for nothing (score 0), + # the same ones every time evaluations = [[next(validation) for _ in range(3)] for _ in range(2)] - assert evaluations[0] == evaluations[1] - assert evaluations[0][:2] == [[{"text": "t0"}, {"text": "t1"}], [{"text": "t2"}]] + assert [[(batch["texts"], batch["score"].item()) for batch in evaluation] for evaluation in evaluations] == [ + [(["t0", "t1"], 1), (["t2"], 1), (["t0", "t1"], 0)]] * 2 + + +def test_build_iterators_refuse_to_evaluate_a_subset(tmp_path): + config = json_config(tmp_path, {"text": [f"t{i}" for i in range(10)]}, max_eval_samples=5) + config.num_workers, config.pin_memory, config.persistent_workers = 0, False, False + config.micro_batch_size, config.eval_micro_batches = 2, 2 + with pytest.raises(ValueError, match="increase validation.eval_iters"): + build_iterators(config, lambda texts: ({"text": text} for text in texts), + lambda examples: {"n": len(examples)}, rank=0, world_size=1) + + +def test_texts_drop_nul_bytes_and_empty_texts(tmp_path): + config = json_config(tmp_path, {"text": ["a\x00b", "", "\x00", "c"]}, max_eval_samples=4) + assert [e["text"] for e in config.texts("validation")] == ["ab", "c"] + config = json_config(tmp_path, {"text": ["a\x00b", "", "\x00", "c"]}, max_eval_samples=0) + assert sorted(islice(train_texts(config, rank=0, world_size=1), 2)) == ["ab", "c"] + + +def test_train_texts_are_split_across_dataloader_workers(tmp_path, monkeypatch): + """Each dataloader worker keeps every num_workers-th text of its rank: together, one epoch.""" + config = json_config(tmp_path, {"text": [f"t{i}" for i in range(10)]}, max_eval_samples=0) + workers = [] + for worker_id in range(2): + info = type("WorkerInfo", (), {"id": worker_id, "num_workers": 2}) + monkeypatch.setattr(torch.utils.data, "get_worker_info", lambda info=info: info) + workers.append(list(islice(train_texts(config, rank=0, world_size=1), 5))) + assert not set(workers[0]) & set(workers[1]) + assert sorted(workers[0] + workers[1]) == sorted(f"t{i}" for i in range(10)) + + +def test_baseline_counts_the_bytes_of_its_texts(): + """Each token counts the text since the previous one: every byte once, also with characters split across tokens.""" + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-70m") + texts = ["Hello wรถrld ๐Ÿ‘‹ ืฉืœื•ื ๆ—ฅๆœฌ่ชž ๐Ÿง‘โ€๐Ÿคโ€๐Ÿง‘", "a"] + (ids, _, nbytes), = data_utils.token_chunks(texts, tokenizer, length=64) + assert sum(nbytes) == sum(len(text.encode()) for text in texts) + assert [n for i, n in zip(ids, nbytes, strict=True) if i == tokenizer.eos_token_id] == [0] * ids.count( + tokenizer.eos_token_id) + + +@pytest.mark.parametrize("seq_length", [4, 5, 6, 8]) +def test_baseline_scores_the_bytes_welt_scores(seq_length): + """Every token is a label once (chunks overlap by one token); shift blocks are not scored, as in WeLT.""" + from transformers import AutoTokenizer + + from welt.processor import TextImageProcessor, next_word_labels + + tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-70m") + texts = ["\x0eHello world\x0f ืฉืœื•ื ืขื•ืœื", "no shift block here", "\x0eb\x0f c"] + examples = list(data_utils.token_examples(texts, tokenizer, length=4)) + labels = [token for example in examples for token in example["labels"][example["loss_mask"].bool()].tolist()] + # The labels: each document (after the first, after its EOS) without the text of its shift block + assert tokenizer.decode(labels) == tokenizer.eos_token.join(data_utils.SHIFT_BLOCK.sub("\x0e", t) for t in texts) + baseline_bytes = sum(int(example["label_bytes"].sum()) for example in examples) + + processor = TextImageProcessor.create(max_word_length=32, render_images=False) + welt_bytes = 0 + # Packed in short examples: documents are split across examples, also within shift blocks, as long documents are + # in training + for packed in pack_words(map(processor.pretokenize, texts), seq_length=seq_length): + 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()) + assert baseline_bytes == welt_bytes == sum(len(data_utils.SHIFT_BLOCK.sub("\x0e", t).encode()) for t in texts) diff --git a/tests/test_inference.py b/tests/test_inference.py new file mode 100644 index 0000000..695b412 --- /dev/null +++ b/tests/test_inference.py @@ -0,0 +1,44 @@ +import pytest +import torch + +pytest.importorskip("vllm") + +from welt import ContextTooLongError # noqa: E402 +from welt.inference import WeLTGenerator # noqa: E402 +from welt.processor import TextImageProcessor # noqa: E402 + + +def stub_generator(next_words, max_model_len=100): + """A WeLTGenerator without engines: _next_words(prefixes) returns each active text's next word (bytes).""" + generator = object.__new__(WeLTGenerator) + generator.processor = TextImageProcessor.create(max_word_length=32, render_images=False) + model_config = type("ModelConfig", (), {"max_model_len": max_model_len}) + generator.latent = type("LLM", (), {"llm_engine": type("Engine", (), {"model_config": model_config})}) + generator._latents = lambda sequences: torch.zeros(len(sequences), 1) + calls = [] + + def _next_words(latents, prefixes, sampling): + calls.append((list(prefixes), [params.seed for params in sampling])) + return [next_words(prefix) for prefix in prefixes] + + generator._next_words = _next_words + return generator, calls + + +def test_generation_continues_a_partial_word_and_stops_on_an_empty_word(): + words = iter([b"world ", b""]) + generator, calls = stub_generator(lambda prefix: prefix + b"ld " if prefix else next(words)) + assert generator.generate_words(["hello wor"]) == [["ld ", "world "]] + assert [prefixes for prefixes, _ in calls] == [[b"wor"], [b""], [b""]] + + +def test_generation_stops_at_max_generated_words_with_a_seed_per_text(): + generator, calls = stub_generator(lambda prefix: b"x ") + assert generator.generate(["a ", "b "], max_generated_words=2, seed=10) == ["x x ", "x x "] + assert [seeds for _, seeds in calls] == [[10, 11], [12, 13]] # Distinct across texts and steps + + +def test_generation_refuses_prompts_beyond_the_context(): + generator, _ = stub_generator(lambda prefix: b"x ", max_model_len=4) + with pytest.raises(ContextTooLongError): + generator.generate(["a b c"], max_generated_words=2) diff --git a/tests/test_metrics.py b/tests/test_metrics.py index c301c0a..7fdd175 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -25,46 +25,6 @@ def test_welt_loss_func_metrics(): assert report["word accuracy"].tolist() == [2, 3] -def test_baseline_counts_the_bytes_of_its_texts(): - """Each token counts the text since the previous one: every byte once, also with characters split across tokens.""" - from transformers import AutoTokenizer - - tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-70m") - texts = ["Hello wรถrld ๐Ÿ‘‹ ืฉืœื•ื ๆ—ฅๆœฌ่ชž ๐Ÿง‘โ€๐Ÿคโ€๐Ÿง‘", "a"] - (ids, _, nbytes), = baseline.token_chunks(texts, tokenizer, length=64) - assert sum(nbytes) == sum(len(text.encode()) for text in texts) - assert [n for i, n in zip(ids, nbytes, strict=True) if i == tokenizer.eos_token_id] == [0] * ids.count( - tokenizer.eos_token_id) - - -def test_baseline_scores_the_bytes_welt_scores(): - """Every token is a label once (chunks overlap by one token); shift blocks are not scored, as in WeLT.""" - from transformers import AutoTokenizer - - from welt.processor import TextImageProcessor, next_word_labels - - tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-70m") - texts = ["\x0eHello world\x0f ืฉืœื•ื ืขื•ืœื", "no shift block here", "\x0eb\x0f c"] - examples = list(baseline.token_examples(texts, tokenizer, length=4)) - labels = [token for example in examples for token in example["labels"][example["loss_mask"].bool()].tolist()] - # The labels: each document (after the first, after its EOS) without the text of its shift block - 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 - # 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()) - assert baseline_bytes == welt_bytes == sum(len(baseline.SHIFT_BLOCK.sub("\x0e", t).encode()) for t in texts) - - def test_baseline_loss_func_bits_per_byte(): losses = torch.tensor([[2.0, 3.0, 1.0, 1.5, 9.0]]) loss_mask = torch.tensor([[1, 1, 1, 1, 0]]) # The last label is padding diff --git a/tests/test_processor.py b/tests/test_processor.py index 2e2cf46..ba43099 100644 --- a/tests/test_processor.py +++ b/tests/test_processor.py @@ -1,4 +1,3 @@ -import contextlib import pickle import tempfile @@ -131,6 +130,9 @@ def test_processor_save_and_load_works(processor): assert new_processor.renderer is not None assert new_processor.max_word_length == processor.max_word_length assert new_processor.pretokenize("hello world") == processor.pretokenize("hello world") + # Words longer than max_word_length are split alike (words-segmentation does not save max_bytes) + long_text = "internationalization " * 3 + assert new_processor.pretokenize(long_text) == processor.pretokenize(long_text) def test_processor_save_and_load_works_without_renderer(text_processor): @@ -163,11 +165,11 @@ def test_multiple_shift_blocks(): def test_shift_blocks_do_not_span_packed_sequences(processor): + """A block cut between two sequences is two: to the end of the first, and from the start of the second.""" words = ["a", SO, "b", SI, "c"] - with pytest.warns(UserWarning, match="ShiftOut without ShiftIn"): - result = processor.process_single_example(words, [3, 2]) - assert result["block_ids"].tolist() == [0, 0, 0, 0, 0] - assert result["label_mask"].all() + result = processor.process_single_example(words, [3, 2]) + assert result["block_ids"].tolist() == [0, 1, 1, 2, 0] + assert result["label_mask"].tolist() == [True, False, True, True, True] def test_bpe_pretokenizer_words_are_text_spans(): @@ -181,6 +183,16 @@ def test_bpe_pretokenizer_words_are_text_spans(): assert ControlTokens.ShiftIn in words +def test_bpe_pretokenizer_splits_long_tokens(): + """Tokens longer than max_word_length (with BOS and EOS), e.g. Pythia's runs of spaces, are split, not truncated.""" + processor = TextImageProcessor.create(max_word_length=8, render_images=False, + pretokenizer_name="EleutherAI/pythia-14m") + text = "a" + " " * 20 + "ืฉืœื•ืืฉืœื•ื" + words = processor.pretokenize(text) + assert "".join(words) == ControlTokens.StartOfText + text + assert all(len(word.encode()) <= 6 for word in words) + + def test_collate_fn_pads_every_dimension_and_keeps_dtypes(): batch = [{"mask": torch.tensor([[True, False]]), "ids": torch.tensor([1])}, {"mask": torch.tensor([[True], [True]]), "ids": torch.tensor([2, 3])}] @@ -191,7 +203,7 @@ def test_collate_fn_pads_every_dimension_and_keeps_dtypes(): assert torch.equal(collated["ids"], torch.tensor([[1, 0], [2, 3]])) -@pytest.mark.parametrize("seq_length", [16, 40, 128]) +@pytest.mark.parametrize("seq_length", [5, 8, 16, 40, 128]) # Short: documents split across examples 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.""" documents = [text_processor.pretokenize(text) for text in ORACLE_TEXTS * 2] @@ -201,6 +213,8 @@ def test_labels_match_an_oracle_of_the_words(text_processor, seq_length): 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] + if not example["label_mask"].any(): # A chunk within a shift block + continue 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) @@ -210,16 +224,15 @@ def test_labels_match_an_oracle_of_the_words(text_processor, seq_length): if label is not None] -@pytest.mark.parametrize(("words", "warning", "blocks"), [ - (["a", "b", "c"], None, []), - (["a", SO, "b", SI, "c"], None, [(1, 3)]), - ([SO, SI, "c"], None, [(0, 1)]), - (["a", SO, "b", SI, "c", SO, "d", SI], None, [(1, 3), (5, 7)]), - (["a", SO, "b"], "ShiftOut without ShiftIn", []), - (["a", SI, "b"], "ShiftIn without ShiftOut", []), - (["a", SO, "b", SO, "c", SI], "nested shift blocks", [(3, 5)]), +@pytest.mark.parametrize(("words", "blocks"), [ + (["a", "b", "c"], []), + (["a", SO, "b", SI, "c"], [(1, 3)]), + ([SO, SI, "c"], [(0, 1)]), + (["a", SO, "b", SI, "c", SO, "d", SI], [(1, 3), (5, 7)]), + (["a", SO, "b"], [(1, 2)]), # Cut at the end of a sequence: to its last word + (["a", SI, "b"], [(0, 1)]), # Cut at its start: from its first word + (["a", SO, "b", SO, "c", SI], [(1, 5)]), # A ShiftOut within a block belongs to it (as the baseline's regex) ]) -def test_shift_blocks(words, warning, blocks): +def test_shift_blocks(words, blocks): """Shift blocks span ShiftOut to ShiftIn, inclusive.""" - with pytest.warns(UserWarning, match=warning) if warning else contextlib.nullcontext(): - assert list(get_shift_blocks(words)) == blocks + assert list(get_shift_blocks(words)) == blocks diff --git a/tests/test_server.py b/tests/test_server.py index 2b4a7fa..5c3a9de 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -6,20 +6,21 @@ from werkzeug.serving import make_server # noqa: E402 +from welt import ContextTooLongError # noqa: E402 from welt.server import MODEL_TAG_HEADER, create_app, generate # noqa: E402 class FakeGenerator: def __init__(self): - self.release = threading.Event() + self.entered, self.release = threading.Event(), threading.Event() self.release.set() - self.processor = type("Processor", (), {"pretokenize": staticmethod(lambda text: ["\x02", *text.split()])}) - def generate(self, texts, max_generated_words, temperature, seed): + def generate_words(self, texts, max_generated_words, temperature, seed): if max_generated_words > 100: - raise ValueError("Prompt words + max_generated_words exceed the latent's context") + raise ContextTooLongError("Prompt words + max_generated_words exceed the latent's context") + self.entered.set() self.release.wait() - return [f"{text}!" * max_generated_words for text in texts] + return [[f"{text}!"] * max_generated_words for text in texts] @pytest.fixture @@ -55,10 +56,8 @@ def test_busy_server_answers_503(client, generator): first = threading.Thread(target=client.post, args=("/generate",), kwargs={"json": {"texts": ["a"]}}) first.start() try: - for _ in range(100): # Until the first request holds the generator - response = client.post("/generate", json={"texts": ["b"]}) - if response.status_code == 503: - break + assert generator.entered.wait(timeout=10) # The first request holds the generator + response = client.post("/generate", json={"texts": ["b"]}) assert response.status_code == 503 assert response.headers["Retry-After"] == "1" finally: diff --git a/tests/test_train.py b/tests/test_train.py index f17d1ab..ee82a0c 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -2,6 +2,7 @@ import os import subprocess import sys +from pathlib import Path import pytest import torch @@ -17,7 +18,7 @@ def run(*args): # A port of its own: the in-process Megatron test fixture holds MASTER_PORT - env = {**os.environ, "PYTHONPATH": os.getcwd(), "MASTER_PORT": str(free_port())} + env = {**os.environ, "PYTHONPATH": str(Path(__file__).resolve().parents[1]), "MASTER_PORT": str(free_port())} subprocess.run([*args], check=True, env=env) diff --git a/welt/__init__.py b/welt/__init__.py index e69de29..ee73c19 100644 --- a/welt/__init__.py +++ b/welt/__init__.py @@ -0,0 +1,2 @@ +class ContextTooLongError(ValueError): + """A prompt (and the words to generate) longer than the latent transformer's context: the client's error.""" diff --git a/welt/inference.py b/welt/inference.py index 23fc174..509ec74 100644 --- a/welt/inference.py +++ b/welt/inference.py @@ -26,6 +26,7 @@ from vllm.config import PoolerConfig # noqa: E402 from words_segmentation.pretokenizer import is_word_complete # noqa: E402 +from welt import ContextTooLongError # noqa: E402 from welt.patches import PatchEmbedding # noqa: E402 from welt.processor import TextImageProcessor, get_shift_blocks # noqa: E402 from welt.utf8 import UTF8LogitsProcessor # noqa: E402 @@ -131,11 +132,17 @@ def _next_words(self, latents: torch.Tensor, prefixes: list[bytes], def generate(self, texts: list[str], max_generated_words: int = 50, temperature: float = 0.0, seed: int | None = None) -> list[str]: - """Greedy (or sampled, with temperature) continuation of each text, word by word.""" + """Greedy (or sampled, with temperature) continuation of each text.""" + return ["".join(words) for words in self.generate_words(texts, max_generated_words, temperature, seed)] + + def generate_words(self, texts: list[str], max_generated_words: int = 50, temperature: float = 0.0, + seed: int | None = None) -> list[list[str]]: + """The words of each text's continuation, generated word by word (a text's own seed, when sampling).""" sequences = [self.processor.pretokenize(text) for text in texts] max_length = self.latent.llm_engine.model_config.max_model_len if any(len(words) + max_generated_words > max_length for words in sequences): - raise ValueError(f"Prompt words + max_generated_words exceed the latent's context ({max_length} words)") + raise ContextTooLongError( + f"Prompt words + max_generated_words exceed the latent's context ({max_length} words)") # A prompt ending mid-word continues that word: its bytes become the decoder's prefix max_prefix = self.processor.max_word_length - 3 # Words have at most max_word_length - 2 bytes; one is new prefixes = [b""] * len(texts) @@ -149,7 +156,8 @@ def generate(self, texts: list[str], max_generated_words: int = 50, temperature: if not active: break # Words have at most max_word_length - 2 bytes (without BOS and EOS), as in training - sampling = [SamplingParams(temperature=temperature, seed=None if seed is None else seed + step, + sampling = [SamplingParams(temperature=temperature, + seed=None if seed is None else seed + step * len(texts) + i, max_tokens=self.processor.max_word_length - 2 - len(prefixes[i]), stop_token_ids=[EOS_TOKEN_ID], detokenize=False) for i in active] @@ -165,4 +173,4 @@ def generate(self, texts: list[str], max_generated_words: int = 50, temperature: sequences[i].append(word.decode("utf-8", errors="replace")) still_active.append(i) active = still_active - return ["".join(words) for words in generated] + return generated diff --git a/welt/processor.py b/welt/processor.py index 13fada0..b25036a 100644 --- a/welt/processor.py +++ b/welt/processor.py @@ -1,7 +1,6 @@ import itertools import json import os -import warnings import torch from cachetools import LRUCache @@ -37,31 +36,41 @@ def next_word_labels(input_ids: torch.Tensor, sequence_ids: torch.Tensor, label_ def get_shift_blocks(words: list[str]): - """(start, end) of each shift block: a ShiftOut word, up to the next ShiftIn word (inclusive). Warns about - unmatched ShiftOut and ShiftIn words, which start or end no block.""" + """(start, end) of each shift block: a ShiftOut word, up to the next ShiftIn word (inclusive), as the baseline's + SHIFT_BLOCK finds them in a whole document (a ShiftOut within a block belongs to it). A sequence cut from a longer + document (pack_words) can end within a block, which then runs to its last word (which, continuing in the next + sequence, predicts nothing), or begin within one, which then starts at its first word.""" shift_out = None for i, word in enumerate(words): if word == ControlTokens.ShiftOut: - if shift_out is not None: - warnings.warn("ShiftOut after ShiftOut without ShiftIn: nested shift blocks are not allowed", - stacklevel=2) - shift_out = i - elif word == ControlTokens.ShiftIn: if shift_out is None: - warnings.warn("ShiftIn without ShiftOut: skipping the shift block", stacklevel=2) - else: - yield shift_out, i - shift_out = None + shift_out = i + elif word == ControlTokens.ShiftIn: + yield (0 if shift_out is None else shift_out), i + shift_out = None if shift_out is not None: - warnings.warn("ShiftOut without ShiftIn at the end of the sequence", stacklevel=2) + yield shift_out, len(words) - 1 + + +def split_bytes(text: str, max_bytes: int) -> list[str]: + """Split a text into pieces of at most max_bytes UTF-8 bytes, between characters.""" + pieces, piece, size = [], "", 0 + for char in text: + length = len(char.encode()) + if piece and size + length > max_bytes: + pieces.append(piece) + piece, size = "", 0 + piece, size = piece + char, size + length + return [*pieces, piece] if piece else pieces class TextImageProcessor: """Turns text into word-level byte tensors, rendered word patches, and labels.""" def __init__(self, pretokenizer: PreTrainedTokenizer, renderer: PixelRendererProcessor | None, - max_word_length: int = 32): + max_word_length: int = 32, trust_remote_code: bool = False): self.pretokenizer = pretokenizer + self.trust_remote_code = trust_remote_code self.tokenizer = UTF8Tokenizer() self.renderer = renderer self.max_word_length = max_word_length @@ -76,14 +85,15 @@ def create(cls, max_word_length: int, render_images: bool, else: pretokenizer = WordsSegmentationTokenizer(max_bytes=max_word_length - 2) # BOS and EOS renderer = PixelRendererProcessor(font=FontConfig(sources=FONTS_NOTO_SANS)) if render_images else None - return cls(pretokenizer=pretokenizer, renderer=renderer, max_word_length=max_word_length) + return cls(pretokenizer=pretokenizer, renderer=renderer, max_word_length=max_word_length, + trust_remote_code=trust_remote_code) def save_pretrained(self, save_directory): os.makedirs(save_directory, exist_ok=True) self.pretokenizer.save_pretrained(os.path.join(save_directory, "pretokenizer")) if self.renderer is not None: self.renderer.save_pretrained(os.path.join(save_directory, "renderer")) - config = {"max_word_length": self.max_word_length} + config = {"max_word_length": self.max_word_length, "trust_remote_code": self.trust_remote_code} with open(os.path.join(save_directory, PROCESSOR_CONFIG_NAME), "w") as f: json.dump(config, f, indent=2, sort_keys=True) @@ -93,8 +103,13 @@ def from_pretrained(cls, path): config = json.load(f) renderer_dir = os.path.join(path, "renderer") renderer = PixelRendererProcessor.from_pretrained(renderer_dir) if os.path.isdir(renderer_dir) else None - return cls(pretokenizer=AutoTokenizer.from_pretrained(os.path.join(path, "pretokenizer")), renderer=renderer, - max_word_length=config["max_word_length"]) + trust_remote_code = config.get("trust_remote_code", False) + pretokenizer = AutoTokenizer.from_pretrained(os.path.join(path, "pretokenizer"), + trust_remote_code=trust_remote_code) + if isinstance(pretokenizer, WordsSegmentationTokenizer): # Its saved config lacks max_bytes + pretokenizer.max_bytes = config["max_word_length"] - 2 + return cls(pretokenizer=pretokenizer, renderer=renderer, max_word_length=config["max_word_length"], + trust_remote_code=trust_remote_code) def render_texts(self, texts: list[str]) -> tuple[torch.Tensor, torch.Tensor]: """Render words into their 16x16 patches, packed: (total patches, 768) uint8, each word's row-major patches in @@ -115,10 +130,12 @@ def pretokenize(self, text: str) -> list[str]: text = self.tokenizer.bos_token + text if isinstance(self.pretokenizer, WordsSegmentationTokenizer): return self.pretokenizer.tokenize(text) - # Other tokenizers' tokens can be encoded (e.g. byte-level BPE "ฤ world"), words are spans of the text instead + # Other tokenizers' tokens can be encoded (e.g. byte-level BPE "ฤ world"), words are spans of the text instead, + # split (as words-segmentation's words) to fit max_word_length with BOS and EOS, rather than truncated offsets = self.pretokenizer(text, add_special_tokens=False, return_offsets_mapping=True).offset_mapping starts = sorted({0, *(start for start, _ in offsets)}) - return [text[start:end] for start, end in zip(starts, [*starts[1:], len(text)], strict=True) if end > start] + spans = (text[start:end] for start, end in zip(starts, [*starts[1:], len(text)], strict=True) if end > start) + return [piece for span in spans for piece in split_bytes(span, self.max_word_length - 2)] def tokenize_words(self, words: list[str]): return self.tokenizer.torch( diff --git a/welt/server.py b/welt/server.py index f697287..99ec45f 100644 --- a/welt/server.py +++ b/welt/server.py @@ -21,6 +21,8 @@ from flask import Flask, jsonify, request +from welt import ContextTooLongError + MODEL_TAG_HEADER = "X-Model-Tag" # Which model revision served a response, from the MODEL_TAG environment variable EXPORT_DIR_ENV = "WELT_EXPORT_DIR" @@ -60,13 +62,13 @@ def generate_route(): if not busy.acquire(blocking=False): return jsonify(message="Busy"), 503, {"Retry-After": "1"} try: - outputs = generator.generate(texts, **options) - except ValueError as error: # e.g. a context longer than the model's + words = generator.generate_words(texts, **options) + except ContextTooLongError as error: return jsonify(message=str(error)), 400 finally: busy.release() - generated_words = sum(len(generator.processor.pretokenize(output)) - 1 for output in outputs) # Without BOS - return jsonify(outputs=outputs, generated_words=generated_words) + return jsonify(outputs=["".join(text_words) for text_words in words], + generated_words=sum(len(text_words) for text_words in words)) return app @@ -94,8 +96,9 @@ def generate_main(): parser.add_argument("--url", default="http://localhost:8080") parser.add_argument("--max_generated_words", type=int, default=50) parser.add_argument("--temperature", type=float, default=0.0) + parser.add_argument("--seed", type=int, default=None) args = parser.parse_args() - outputs = generate(args.url, args.prompts, args.max_generated_words, args.temperature)["outputs"] + outputs = generate(args.url, args.prompts, args.max_generated_words, args.temperature, args.seed)["outputs"] for prompt, output in zip(args.prompts, outputs, strict=True): print(repr(prompt), "->", repr(output)) diff --git a/welt_training/README.md b/welt_training/README.md index 0660647..9a47c45 100644 --- a/welt_training/README.md +++ b/welt_training/README.md @@ -87,9 +87,10 @@ encoders and the bytes decoder only compute on real bytes. Resuming from a check ### Validation -Each evaluation draws `validation.eval_iters` global batches. The validation examples (split across ranks) are -repeated to fill them, the same ones at every evaluation; if they do not fit, evaluations see only the first ones, and -a warning is logged: raise `eval_iters` to cover them. +Each evaluation draws `validation.eval_iters` global batches, and scores every validation example (split across +ranks) exactly once, the same ones at every evaluation: the remaining micro batches are copies that count for nothing. +If the examples do not fit, building the dataset fails: raise `eval_iters` to cover them (two models then score the +same texts, whatever their batch sizes). ## Parallelism diff --git a/welt_training/baseline.py b/welt_training/baseline.py index b34173d..6ac37a2 100644 --- a/welt_training/baseline.py +++ b/welt_training/baseline.py @@ -12,8 +12,6 @@ Like WeLT, it does not predict (or score) the text of shift blocks (`\x0E...\x0F`, e.g. the source sentence). """ import math -import re -from collections.abc import Iterable, Iterator from dataclasses import dataclass from functools import partial @@ -25,47 +23,9 @@ from transformers import AutoTokenizer from welt.model import safetensors_checkpoint, transformer_provider -from welt_training.data_utils import TextDataConfig, build_iterators +from welt_training.data_utils import TextDataConfig, build_iterators, token_examples from welt_training.train import build_config, main, report, to_cuda -SHIFT_BLOCK = re.compile("\x0e[^\x0f]*\x0f") - - -def token_chunks(texts: Iterable[str], tokenizer, length: int) -> Iterator[tuple[list[int], list[int], list[int]]]: - """Tokenize texts, concatenate them (each after an EOS), and split into chunks of `length` tokens: (ids, loss - mask, UTF-8 bytes of each token's text). Chunks overlap by one token, so that every token is a label exactly once; - the last one is padded with EOS. The loss mask is 0 for the padding, and for the tokens within shift blocks - (after \x0E, up to \x0F). A token's bytes are the text since the previous token's end (0 for EOS, and for the - second part of a character split across tokens).""" - ids, mask, nbytes = [], [], [] - for text in texts: - encoded = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True) - blocks = [(match.start(), match.end()) for match in SHIFT_BLOCK.finditer(text)] - ids += [tokenizer.eos_token_id, *encoded.input_ids] - mask += [1, *(int(not any(start + 1 < end <= block_end for start, block_end in blocks)) - for _, end in encoded.offset_mapping)] - nbytes.append(0) - previous = 0 - for _, end in encoded.offset_mapping: - nbytes.append(len(text[previous:max(end, previous)].encode())) - previous = max(end, previous) - while len(ids) >= length: - yield ids[:length], mask[:length], nbytes[:length] - ids, mask, nbytes = ids[length - 1:], mask[length - 1:], nbytes[length - 1:] - if len(ids) > 1: - padding = length - len(ids) - yield ids + [tokenizer.eos_token_id] * padding, mask + [0] * padding, nbytes + [0] * padding - - -def token_examples(texts: Iterable[str], tokenizer, length: int) -> Iterator[dict]: - """Token chunks of length + 1 as inputs, labels, their loss mask, the UTF-8 bytes of each scored label, and which - labels are scored text (not EOS).""" - for ids, loss_mask, nbytes in token_chunks(texts, tokenizer, length + 1): - ids, loss_mask = torch.tensor(ids), torch.tensor(loss_mask[1:]) - yield {"input_ids": ids[:-1], "labels": ids[1:], "loss_mask": loss_mask, - "label_bytes": torch.tensor(nbytes[1:]) * loss_mask, - "text_mask": loss_mask * (ids[1:] != tokenizer.eos_token_id)} - @dataclass(kw_only=True) class TokensDatasetProvider(TextDataConfig, DatasetProvider): @@ -81,8 +41,11 @@ def build_datasets(self, context: DatasetBuildContext): return train, validation, None -def loss_func(losses: torch.Tensor, loss_mask: torch.Tensor, label_bytes: torch.Tensor, text_mask: torch.Tensor): - """Per-token cross entropy, plus bits per byte over the scored text tokens (not EOS).""" +def loss_func(losses: torch.Tensor, loss_mask: torch.Tensor, label_bytes: torch.Tensor, text_mask: torch.Tensor, + score: torch.Tensor | None = None): + """Per-token cross entropy, plus bits per byte over the scored text tokens (not EOS). score: as WeLT's.""" + if score is not None: + loss_mask, label_bytes, text_mask = loss_mask * score, label_bytes * score, text_mask * score losses = losses.float() loss, num_tokens, metrics = masked_next_token_loss(loss_mask, losses) # With NaN / Inf checks return loss, num_tokens, metrics | { @@ -92,11 +55,12 @@ def loss_func(losses: torch.Tensor, loss_mask: torch.Tensor, label_bytes: torch. def forward_step(state, data_iterator, model, return_schedule_plan: bool = False): batch = to_cuda(next(data_iterator)) + score = batch.pop("score", None) input_ids = batch["input_ids"] position_ids = torch.arange(input_ids.size(1), device=input_ids.device).expand_as(input_ids) losses = model(input_ids=input_ids, position_ids=position_ids, attention_mask=None, labels=batch["labels"]) return losses, partial(loss_func, loss_mask=batch["loss_mask"], label_bytes=batch["label_bytes"], - text_mask=batch["text_mask"]) + text_mask=batch["text_mask"], score=score) def build(config: dict): @@ -115,9 +79,15 @@ def build(config: dict): # 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"]) + # Other model options are the transformer's (e.g. hidden_dropout), as WeLT's are its latent transformer's + for key, value in model_config.items(): + if key not in ("transformer", "tokenizer", "load_pretrained", "trust_remote_code"): + if not hasattr(model, key): + raise ValueError(f"Unknown model option {key!r}") + setattr(model, key, value) + + dataset = TokensDatasetProvider(tokenizer_name=model_config["tokenizer"], + **{"trust_remote_code": trust_remote_code} | config["data"]) return build_config(config, model, dataset, vocab_size=model.vocab_size), None diff --git a/welt_training/data.py b/welt_training/data.py index 5195ae7..6a15444 100644 --- a/welt_training/data.py +++ b/welt_training/data.py @@ -1,6 +1,7 @@ """Megatron-Bridge dataset provider for WeLT: streamed texts, split into words and packed into examples on the fly.""" -from collections.abc import Iterable +from collections.abc import Iterable, Iterator from dataclasses import dataclass +from functools import partial from megatron.bridge.training.config import DatasetBuildContext, DatasetProvider from megatron.core import parallel_state @@ -23,13 +24,14 @@ def processor(self) -> TextImageProcessor: render_images=self.render_images, pretokenizer_name=self.pretokenizer_name, trust_remote_code=self.trust_remote_code) - def build_datasets(self, context: DatasetBuildContext): - processor = self.processor() - - def make_examples(texts: Iterable[str]): - for packed in pack_words(map(processor.pretokenize, texts), self.seq_length): - yield processor.process_single_example(*packed) + def make_examples(self, texts: Iterable[str], processor: TextImageProcessor | None = None) -> Iterator[dict]: + """Packed examples of the texts' words (see pack_words), as in training.""" + processor = processor or self.processor() + for packed in pack_words(map(processor.pretokenize, texts), self.seq_length): + yield processor.process_single_example(*packed) + def build_datasets(self, context: DatasetBuildContext): + make_examples = partial(self.make_examples, processor=self.processor()) train, validation = build_iterators(self, make_examples, collate_fn, rank=parallel_state.get_data_parallel_rank(), world_size=parallel_state.get_data_parallel_world_size()) diff --git a/welt_training/data_utils.py b/welt_training/data_utils.py index a8168a7..79480b3 100644 --- a/welt_training/data_utils.py +++ b/welt_training/data_utils.py @@ -2,6 +2,7 @@ ranks and dataloader workers. Examples are made from the texts on the fly, in the dataloader workers. Independent of Megatron.""" import logging +import re from collections.abc import Callable, Iterable, Iterator from dataclasses import dataclass from itertools import count, cycle, islice @@ -36,8 +37,8 @@ class TextDataConfig: eval_micro_batches: int = 1 micro_batch_size: int = 1 - def examples(self, split: str) -> IterableDataset: - """The raw examples of the train or validation split, streamed.""" + def examples(self, split: str, rank: int = 0, world_size: int = 1) -> IterableDataset: + """The raw examples of the train or validation split, streamed; train examples of a data parallel rank.""" has_validation = "validation" in get_dataset_split_names(self.dataset_name, self.dataset_config_name, data_files=self.data_files) dataset = load_dataset(self.dataset_name, self.dataset_config_name, data_files=self.data_files, @@ -45,44 +46,55 @@ def examples(self, split: str) -> IterableDataset: streaming=True) if split == "validation": return dataset.take(self.max_eval_samples) + # Split before skipping the held out validation examples: each is among the first of its rank's stream + # (whether ranks get whole files or every world_size-th example) + dataset = split_dataset_by_node(dataset, rank=rank, world_size=world_size) return dataset if has_validation else dataset.skip(self.max_eval_samples) - def texts(self, split: str) -> IterableDataset: - """The formatted (non-empty) texts of a split, streamed, as {"text": ...} examples.""" + def text_formatter(self) -> Callable[[dict], str]: + """A raw example's text: decoded (dataset_tokenizer), formatted (dataset_text_template), without NULs.""" 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): + def format_text(example: dict) -> str: 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) - return {"text": text.replace("\x00", "")} # NUL is the PAD byte + return text.replace("\x00", "") # NUL is the PAD byte + return format_text + + def texts(self, split: str) -> IterableDataset: + """The formatted (non-empty) texts of a split, streamed, as {"text": ...} examples.""" + format_text = self.text_formatter() dataset = self.examples(split) - return dataset.map(format_text, remove_columns=dataset.column_names).filter(lambda e: len(e["text"]) > 0) + return (dataset.map(lambda example: {"text": format_text(example)}, remove_columns=dataset.column_names) + .filter(lambda example: len(example["text"]) > 0)) def train_texts(config: TextDataConfig, rank: int, world_size: int) -> Iterator[str]: """Endless shuffled train texts of this data parallel rank and dataloader worker: each worker reads its rank's - stream (shuffled alike) and keeps every num_workers-th text, so all workers make examples however many files - the dataset has. Epochs reshuffle.""" + stream (shuffled alike), keeps every num_workers-th example, and formats only those, so all workers make examples + however many files the dataset has. Epochs reshuffle.""" # ponytail: resuming from a checkpoint restarts the stream; track the texts consumed if exact resumption matters - dataset = split_dataset_by_node(config.texts("train").shuffle(seed=config.seed, - buffer_size=config.shuffle_buffer_size), - rank=rank, world_size=world_size) + dataset = config.examples("train", rank, world_size).shuffle(seed=config.seed, + buffer_size=config.shuffle_buffer_size) + format_text = config.text_formatter() worker = torch.utils.data.get_worker_info() worker_id, num_workers = (worker.id, worker.num_workers) if worker is not None else (0, 1) for epoch in count(): dataset.set_epoch(epoch) # .iter(), unlike __iter__, does not split the stream across dataloader workers by files - texts = (text for batch in dataset.iter(batch_size=1000) for text in batch["text"]) + examples = (dict(zip(batch, values, strict=True)) for batch in dataset.iter(batch_size=1000) + for values in zip(*batch.values(), strict=True)) empty = True - for text in islice(texts, worker_id, None, num_workers): - empty = False - yield text + for example in islice(examples, worker_id, None, num_workers): + if text := format_text(example): + empty = False + yield text if empty: raise ValueError(f"No train texts for data parallel rank {rank}, dataloader worker {worker_id}") @@ -101,8 +113,9 @@ def __iter__(self): def build_iterators(config, make_examples: Callable[[Iterable[str]], Iterator[dict]], collate_fn, rank: int, world_size: int): """Megatron's (external) train and validation iterators of micro batches, for a DatasetProvider with dataloader - options (num_workers, pin_memory, persistent_workers). Validation: the examples of the validation texts, split - across ranks, repeated to fill config.eval_micro_batches, the same ones at every evaluation.""" + options (num_workers, pin_memory, persistent_workers). Validation: every example of the validation texts (split + across ranks) once, in config.eval_micro_batches micro batches, the same ones at every evaluation. Batches are + weighted by a "score": 1, or 0 for the copies that fill the remaining micro batches (and count for nothing).""" # The workers fork here, before this process makes (e.g. renders) any example: a fork while another thread holds # a lock (e.g. Pango's) deadlocks the children train = iter(torch.utils.data.DataLoader( @@ -116,12 +129,12 @@ def build_iterators(config, make_examples: Callable[[Iterable[str]], Iterator[di examples = list(make_examples(texts))[rank::world_size] if not examples: raise ValueError(f"No validation examples for data parallel rank {rank} (from {len(texts)} texts)") - batches = [collate_fn(examples[i:i + config.micro_batch_size]) + batches = [collate_fn(examples[i:i + config.micro_batch_size]) | {"score": torch.ones(())} for i in range(0, len(examples), config.micro_batch_size)] if len(batches) > config.eval_micro_batches: - logger.warning(f"Evaluations see {config.eval_micro_batches} of {len(batches)} validation micro batches, " - "increase validation.eval_iters to see them all") - validation = [batches[i % len(batches)] for i in range(config.eval_micro_batches)] + raise ValueError(f"{len(batches)} validation micro batches, but evaluations see {config.eval_micro_batches}: " + "increase validation.eval_iters to see them all") + validation = batches + [batches[0] | {"score": torch.zeros(())}] * (config.eval_micro_batches - len(batches)) return train, cycle(validation) @@ -152,3 +165,42 @@ def padded(): start += seq_length - 1 if words: yield padded() + + +SHIFT_BLOCK = re.compile("\x0e[^\x0f]*\x0f") + + +def token_chunks(texts: Iterable[str], tokenizer, length: int) -> Iterator[tuple[list[int], list[int], list[int]]]: + """Tokenize texts, concatenate them (each after an EOS), and split into chunks of `length` tokens: (ids, loss + mask, UTF-8 bytes of each token's text). Chunks overlap by one token, so that every token is a label exactly once; + the last one is padded with EOS. The loss mask is 0 for the padding, and for the tokens within shift blocks + (after \x0E, up to \x0F). A token's bytes are the text since the previous token's end (0 for EOS, and for the + second part of a character split across tokens).""" + ids, mask, nbytes = [], [], [] + for text in texts: + encoded = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True) + blocks = [(match.start(), match.end()) for match in SHIFT_BLOCK.finditer(text)] + ids += [tokenizer.eos_token_id, *encoded.input_ids] + mask += [1, *(int(not any(start + 1 < end <= block_end for start, block_end in blocks)) + for _, end in encoded.offset_mapping)] + nbytes.append(0) + previous = 0 + for _, end in encoded.offset_mapping: + nbytes.append(len(text[previous:max(end, previous)].encode())) + previous = max(end, previous) + while len(ids) >= length: + yield ids[:length], mask[:length], nbytes[:length] + ids, mask, nbytes = ids[length - 1:], mask[length - 1:], nbytes[length - 1:] + if len(ids) > 1: + padding = length - len(ids) + yield ids + [tokenizer.eos_token_id] * padding, mask + [0] * padding, nbytes + [0] * padding + + +def token_examples(texts: Iterable[str], tokenizer, length: int) -> Iterator[dict]: + """Token chunks of length + 1 as inputs, labels, their loss mask, the UTF-8 bytes of each scored label, and which + labels are scored text (not EOS).""" + for ids, loss_mask, nbytes in token_chunks(texts, tokenizer, length + 1): + ids, loss_mask = torch.tensor(ids), torch.tensor(loss_mask[1:]) + yield {"input_ids": ids[:-1], "labels": ids[1:], "loss_mask": loss_mask, + "label_bytes": torch.tensor(nbytes[1:]) * loss_mask, + "text_mask": loss_mask * (ids[1:] != tokenizer.eos_token_id)} diff --git a/welt_training/experiments/language-modeling/README.md b/welt_training/experiments/language-modeling/README.md index b676bb6..f7fdf81 100644 --- a/welt_training/experiments/language-modeling/README.md +++ b/welt_training/experiments/language-modeling/README.md @@ -31,8 +31,9 @@ an image encoder: | 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. +Examples hold the same text (1.33 Pythia tokens per word). The first 32 documents are held out for validation, and +every evaluation scores all of them (`validation.eval_iters: 2` covers each model's 72-80 examples), so all models +score 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). @@ -70,7 +71,8 @@ Validation bits per byte after 10,000 steps (the best per latent transformer in - 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. + 1.034 with 54M; `bytes-31m-half.json` is 6.8M, and `bytes-dec-70m.json` a 5.9M decoder for a 1.1M encoder at + matched FLOPs), 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. diff --git a/welt_training/experiments/language-modeling/climbmix-baseline.yaml b/welt_training/experiments/language-modeling/climbmix-baseline.yaml index abd7dae..fe3af18 100644 --- a/welt_training/experiments/language-modeling/climbmix-baseline.yaml +++ b/welt_training/experiments/language-modeling/climbmix-baseline.yaml @@ -13,7 +13,7 @@ data: validation: eval_interval: 250 - eval_iters: 1 + eval_iters: 2 logger: wandb_project: welt-language-modeling diff --git a/welt_training/experiments/language-modeling/climbmix-welt.yaml b/welt_training/experiments/language-modeling/climbmix-welt.yaml index d1006ec..d521fb0 100644 --- a/welt_training/experiments/language-modeling/climbmix-welt.yaml +++ b/welt_training/experiments/language-modeling/climbmix-welt.yaml @@ -13,13 +13,13 @@ data: 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 + # The first documents are held out for validation, all of them scored at every evaluation (the same bytes for all + # models): 80 examples of WeLT (72 of the baseline), within eval_iters global batches max_eval_samples: 32 validation: eval_interval: 250 - eval_iters: 1 + eval_iters: 2 logger: wandb_project: welt-language-modeling diff --git a/welt_training/train.py b/welt_training/train.py index 7ef1a02..d287d61 100644 --- a/welt_training/train.py +++ b/welt_training/train.py @@ -107,11 +107,14 @@ def to_cuda(batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: return {key: value.cuda(non_blocking=True) for key, value in batch.items()} -def loss_func(losses: torch.Tensor, correct: torch.Tensor, labels: torch.Tensor): +def loss_func(losses: torch.Tensor, correct: torch.Tensor, labels: torch.Tensor, score: torch.Tensor | None = None): """Per-byte cross entropy, plus bits per byte and byte/word accuracy for logging. labels: (words, bytes). Bits per byte count every prediction (including the EOS ending each word) except the - EOS ending a document (a word with an empty label), per UTF-8 byte of text, like the causal LM baseline.""" + EOS ending a document (a word with an empty label), per UTF-8 byte of text, like the causal LM baseline. + score: 0 for validation micro batches that only fill an evaluation (build_iterators), which count for nothing.""" loss_mask = labels != PAD_TOKEN_ID + if score is not None: + loss_mask &= score.bool() eos = labels == EOS_TOKEN_ID document_end = torch.zeros_like(eos) document_end[:, 0] = eos[:, 0] # A word whose first label byte is EOS is empty: the end of a document @@ -129,8 +132,9 @@ def loss_func(losses: torch.Tensor, correct: torch.Tensor, labels: torch.Tensor) def forward_step(state, data_iterator, model, return_schedule_plan: bool = False): batch = to_cuda(next(data_iterator)) + score = batch.pop("score", None) losses, correct, labels = model(**batch) - return losses, partial(loss_func, correct=correct, labels=labels) + return losses, partial(loss_func, correct=correct, labels=labels, score=score) def main(build, forward_step_func):