from __future__ import annotations import json import os from pathlib import Path import subprocess import sys import types import importlib.util import pytest import dlb.recipes as recipes_module from dlb.recipes import ( BACKBONE_TENSOR_KEYS, RecipeError, build_launch, expected_backbone_keys, load_recipe, masked_to_absorbing, uniform_to_absorbing, validate_effective_batch, ) ROOT = Path(__file__).parents[1] def override(command: tuple[str, ...], key: str) -> str: prefix = key + "=" values = [argument[len(prefix) :] for argument in command if argument.startswith(prefix)] assert len(values) == 1, (key, command) return values[0] def test_flm_recipe_matches_paper_v3_reference() -> None: recipe = load_recipe("flm", "lm1b") assert recipe.max_steps == 1_000_000 assert recipe.global_batch_size == 512 assert recipe.learning_rate == pytest.approx(3e-4) assert recipe.warmup_steps == 2_500 assert recipe.optimizer == "adam" assert recipe.sequence_length == 128 assert recipe.evidence == "paper_v3_main" def test_fmlm_recipe_uses_main_text_ce_schedule_not_release_drift() -> None: recipe = load_recipe("fmlm", "lm1b") assert recipe.max_steps == 100_000 assert recipe.global_batch_size == 512 assert recipe.learning_rate == pytest.approx(3e-4) assert recipe.warmup_steps == 2_500 assert recipe.distillation_loss == "cross_entropy" assert recipe.release_script_max_steps == 1_000_000 def test_flm_release_script_drift_is_explicit() -> None: recipe = load_recipe("flm", "owt") assert recipe.max_steps == 1_000_000 assert recipe.release_script_max_steps == 1_500_000 assert recipe.release_evidence == "official_release_script" assert "trainer.max_steps=1500000" in ( ROOT / "upstreams" / "flm" / "scripts" / "train_lm1b_flm.sh" ).read_text() assert "trainer.max_steps=1000000" in ( ROOT / "upstreams" / "flm" / "scripts" / "train_lm1b_fmlm_denoiser.sh" ).read_text() def test_sdtt_and_dcd_round_schedule() -> None: for name in ("mdlm_sdtt", "duo_dcd"): recipe = load_recipe(name, "lm1b") assert recipe.rounds == 8 assert recipe.steps_per_round == 10_000 assert recipe.max_steps == 80_000 assert recipe.global_batch_size == 128 assert recipe.learning_rate == pytest.approx(6e-5) assert recipe.warmup_steps == 2_500 @pytest.mark.parametrize( ("dataset", "cache_name"), [ ("lm1b", "bert-base-uncased.pkl"), ("owt", "gpt2.pkl"), ], ) def test_duo_dcd_uses_repo_integral_cache_not_tokenizer_path( dataset: str, cache_name: str ) -> None: launch = build_launch( load_recipe("duo_dcd", dataset), root=ROOT, source=ROOT / "upstreams" / "duo", output=ROOT / "checkpoints" / "self_trained" / dataset / "duo_dcd", teacher=ROOT / "checkpoints" / "fixture" / "uniform_duo.ckpt", devices=8, nodes=1, per_device_batch_size=16, seed=42, resume=True, ) integral_cache = override(launch.command, "algo.integral_cache_path") curriculum_integral_cache = override( launch.command, "+algo.curriculum.integral_cache_path" ) assert integral_cache == str( (ROOT / "upstreams" / "duo" / "integral" / cache_name).absolute() ) assert curriculum_integral_cache == integral_cache assert "snapshots" not in integral_cache assert override(launch.command, "+algo.curriculum.mode") == "simple" assert override(launch.command, "+algo.curriculum.gamma_min") == "-4" assert override(launch.command, "+algo.curriculum.gamma_max") == "-1" assert Path(launch.command[2]).name == "train_duo_dcd.py" def test_duo_dcd_wrapper_accepts_trainer_base_nll_signature( monkeypatch: pytest.MonkeyPatch, ) -> None: spec = importlib.util.spec_from_file_location( "train_duo_dcd_for_test", ROOT / "adapters" / "train_duo_dcd.py" ) assert spec is not None and spec.loader is not None train_duo_dcd = importlib.util.module_from_spec(spec) spec.loader.exec_module(train_duo_dcd) calls = [] class Distillation: def nll( self, x0, output_tokens, current_accumulation_step=None, train_mode=None, ): calls.append( (x0, output_tokens, current_accumulation_step, train_mode) ) return "loss" algo = type(sys)("algo") algo.Distillation = Distillation monkeypatch.setitem(sys.modules, "algo", algo) train_duo_dcd._patch_distillation_nll() assert Distillation().nll("x0", None, "output", "accum", True) == "loss" assert calls == [("x0", "output", "accum", True)] def test_di4c_sampling_checkpoints_match_reference_selection() -> None: assert load_recipe("duo_di4c", "lm1b").sampling_step == 20_000 assert load_recipe("mdlm_di4c", "lm1b").sampling_step == 20_000 assert load_recipe("duo_di4c", "owt").sampling_step == 50_000 def test_lm1b_di4c_uses_project_wrapper_and_teacher_student_init() -> None: recipe = load_recipe("duo_di4c", "lm1b") launch = build_launch( recipe, root=ROOT, source=ROOT / "upstreams" / "di4c", output=ROOT / "checkpoints" / "reference_reproduction" / "duo_di4c" / "lm1b", teacher=ROOT / "checkpoints" / "fixture" / "uniform_duo.ckpt", devices=2, nodes=1, per_device_batch_size=1, seed=42, resume=True, ) assert Path(launch.command[2]).name == "train_di4c.py" assert override(launch.command, "+dlb_student_init") == "teacher" def test_di4c_student_init_can_be_overridden_to_scratch() -> None: recipe = load_recipe("mdlm_di4c", "lm1b") launch = build_launch( recipe, root=ROOT, source=ROOT / "upstreams" / "di4c", output=ROOT / "checkpoints" / "reference_reproduction" / "mdlm_di4c" / "lm1b", teacher=ROOT / "checkpoints" / "fixture" / "masked_mdlm.ckpt", devices=2, nodes=1, per_device_batch_size=1, seed=42, resume=True, di4c_student_init="scratch", ) assert Path(launch.command[2]).name == "train_di4c.py" assert override(launch.command, "+dlb_student_init") == "scratch" def test_di4c_wrapper_registers_hydra_resolvers_before_composition( monkeypatch: pytest.MonkeyPatch, ) -> None: from adapters import train_di4c registered = {} class FakeOmegaConf: @staticmethod def register_new_resolver(name, resolver, *, replace=False): registered[name] = (resolver, replace) monkeypatch.setitem( sys.modules, "omegaconf", types.SimpleNamespace(OmegaConf=FakeOmegaConf), ) train_di4c._register_omegaconf_resolvers() assert set(registered) == {"cwd", "device_count", "div_up", "eval"} assert all(replace for _, replace in registered.values()) assert registered["div_up"][0](5, 2) == 3 assert registered["eval"][0]("1 + 1") == 2 assert registered["eval"][0]("len([1, 2, 3])") == 3 assert registered["eval"][0]("__import__('math').ceil(1.2)") == 2 assert isinstance(registered["cwd"][0](), str) assert isinstance(registered["device_count"][0](), int) def test_di4c_wrapper_adds_original_script_directory_to_import_path( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: from adapters import train_di4c source = tmp_path / "upstreams" / "di4c" package_root = source / "sdtt" / "src" / "sdtt" package_root.mkdir(parents=True) (package_root / "main.py").write_text("# fixture\n", encoding="utf-8") monkeypatch.chdir(source) before = list(sys.path) try: train_di4c._install_upstream_path() assert str(package_root) in sys.path assert sys.path.index(str(package_root)) < sys.path.index( str(source / "sdtt" / "src") ) finally: sys.path[:] = before def test_di4c_hf_small_student_init_preserves_configured_teacher( monkeypatch: pytest.MonkeyPatch, ) -> None: from adapters import train_di4c fake_sdtt = types.ModuleType("sdtt") fake_sdtt.load_small_student = lambda **kwargs: { "source": "upstream-default-loader", "kwargs": kwargs, } fake_core = types.ModuleType("sdtt.core") fake_distill = types.ModuleType("sdtt.core.distill") fake_multi_round = types.ModuleType("sdtt.core.distill.multi_round_sdtt") calls = [] class FakeMultiRoundSDTT: @classmethod def from_pretrained( cls, repo, revision, *, config=None, student_as_teacher=True, ): calls.append( { "repo": repo, "revision": revision, "config": config, "student_as_teacher": student_as_teacher, } ) return {"source": "direct-hf-small-loader"} fake_multi_round.MultiRoundSDTT = FakeMultiRoundSDTT monkeypatch.setitem(sys.modules, "sdtt", fake_sdtt) monkeypatch.setitem(sys.modules, "sdtt.core", fake_core) monkeypatch.setitem(sys.modules, "sdtt.core.distill", fake_distill) monkeypatch.setitem( sys.modules, "sdtt.core.distill.multi_round_sdtt", fake_multi_round, ) train_di4c._patch_student_loader() config = {"dlb_student_init": "hf_small"} result = fake_sdtt.load_small_student( loss="kld", round=7, config=config, student_as_teacher=True, ) assert result == {"source": "direct-hf-small-loader"} assert calls == [ { "repo": "jdeschena/sdtt", "revision": "baselines_kld_step_70000", "config": config, "student_as_teacher": False, } ] @pytest.mark.parametrize( ("model", "dataset"), [ ("flm", "lm1b"), ("duo", "lm1b"), ("mdlm", "lm1b"), ("candi", "owt"), ("rdlm", "lm1b"), ("fmlm", "lm1b"), ("duo_dcd", "lm1b"), ("mdlm_sdtt", "lm1b"), ("duo_di4c", "lm1b"), ("mdlm_di4c", "lm1b"), ], ) def test_rendered_commands_use_real_entrypoints_and_absolute_project_paths( model: str, dataset: str ) -> None: recipe = load_recipe(model, dataset) teacher = ROOT / "checkpoints" / "fixture" / "teacher.ckpt" if recipe.teacher_family else None devices = 2 if model.endswith("_di4c") else 8 launch = build_launch( recipe, root=ROOT, source=ROOT / recipe.source_path, output=ROOT / "checkpoints" / "self_trained" / dataset / model, teacher=teacher, devices=devices, nodes=1, per_device_batch_size=recipe.default_per_device_batch_size, seed=42, resume=True, ) assert Path(launch.command[0]).is_absolute() assert Path(launch.entrypoint).is_absolute() assert Path(launch.cwd).is_absolute() assert Path(launch.data_path).is_absolute() assert Path(launch.output).is_absolute() assert all("${" not in argument for argument in launch.command) assert override(launch.command, recipe.output_override) == str(launch.output) assert override(launch.command, recipe.max_steps_override) == str(recipe.max_steps) assert override(launch.command, recipe.learning_rate_override) == str( recipe.learning_rate ) assert launch.effective_global_batch_size == recipe.global_batch_size def test_rendered_hydra_overrides_exist_in_pinned_upstream_configs() -> None: flm = build_launch( load_recipe("flm", "lm1b"), root=ROOT, source=ROOT / "upstreams" / "flm", output=ROOT / "checkpoints" / "self_trained" / "lm1b" / "flm", teacher=None, devices=8, nodes=1, per_device_batch_size=64, seed=42, resume=True, ) assert override(flm.command, "data") == "lm1b-wrap" assert override(flm.command, "algo") == "flm" assert override(flm.command, "model.length") == "128" assert override(flm.command, "lr_scheduler.num_warmup_steps") == "2500" assert override(flm.command, "trainer.accumulate_grad_batches") == "1" assert override(flm.command, "wandb") == "null" def test_all_manual_recipes_render_unique_site_independent_overrides() -> None: models = ( "flm", "duo", "mdlm", "candi", "fmlm", "duo_dcd", "mdlm_sdtt", "duo_di4c", "mdlm_di4c", ) for model in models: for dataset in ("lm1b", "owt"): recipe = load_recipe(model, dataset) devices = 2 if model.endswith("_di4c") else 8 teacher = ( ROOT / "checkpoints" / "fixture" / f"{recipe.teacher_family}.ckpt" if recipe.teacher_family else None ) launch = build_launch( recipe, root=ROOT, source=ROOT / recipe.source_path, output=ROOT / "checkpoints" / "self_trained" / dataset / model, teacher=teacher, devices=devices, nodes=1, per_device_batch_size=recipe.default_per_device_batch_size, seed=42, resume=True, ) keys = [argument.split("=", 1)[0].lstrip("+") for argument in launch.command[3:]] assert len(keys) == len(set(keys)), (model, dataset, keys) rendered = "\n".join(launch.command) assert "/home/" not in rendered assert "/share/" not in rendered assert "REDACTED" not in rendered assert launch.entrypoint.is_file() rdlm = load_recipe("rdlm", "lm1b") launch = build_launch( rdlm, root=ROOT, source=ROOT / rdlm.source_path, output=ROOT / "checkpoints" / "self_trained" / "lm1b" / "rdlm", teacher=None, devices=8, nodes=1, per_device_batch_size=64, seed=42, resume=True, ) keys = [argument.split("=", 1)[0].lstrip("+") for argument in launch.command[3:]] assert len(keys) == len(set(keys)) assert "/home/" not in "\n".join(launch.command) def test_effective_global_batch_must_be_exact_and_divisible() -> None: assert validate_effective_batch(512, devices=8, nodes=1, per_device_batch_size=16) == 4 with pytest.raises(RecipeError, match="divisible"): validate_effective_batch(512, devices=6, nodes=1, per_device_batch_size=32) with pytest.raises(RecipeError, match="positive"): validate_effective_batch(512, devices=0, nodes=1, per_device_batch_size=32) def tensor_state(vocab: int = 3, hidden: int = 2) -> dict[str, object]: return { "vocab_embed.embedding": [ [10 * row + column + 1 for column in range(hidden)] for row in range(vocab) ], "output_layer.linear.weight": [ [100 + 10 * row + column for column in range(hidden)] for row in range(vocab) ], "output_layer.linear.bias": [200 + row for row in range(vocab)], } def test_uniform_teacher_appends_zero_absorbing_rows() -> None: state = tensor_state() adapted = uniform_to_absorbing(state, expected_keys=set(BACKBONE_TENSOR_KEYS)) assert adapted["vocab_embed.embedding"] == [[1, 2], [11, 12], [21, 22], [0, 0]] assert adapted["output_layer.linear.weight"][-1] == [0, 0] assert adapted["output_layer.linear.bias"] == [200, 201, 202, 0] assert state["vocab_embed.embedding"] == [[1, 2], [11, 12], [21, 22]] def test_uniform_teacher_zeroes_existing_absorbing_row() -> None: state = tensor_state(vocab=4) adapted = uniform_to_absorbing( state, target_vocab_size=4, expected_keys=set(BACKBONE_TENSOR_KEYS), ) assert adapted["vocab_embed.embedding"] == [[1, 2], [11, 12], [21, 22], [0, 0]] assert adapted["output_layer.linear.weight"][-1] == [0, 0] assert adapted["output_layer.linear.bias"] == [200, 201, 202, 0] assert state["vocab_embed.embedding"][-1] == [31, 32] def test_masked_teacher_moves_existing_mask_row_to_final_absorbing_state() -> None: state = tensor_state() adapted = masked_to_absorbing( state, source_mask_index=1, target_vocab_size=4, expected_keys=set(BACKBONE_TENSOR_KEYS), ) assert adapted["vocab_embed.embedding"] == [[1, 2], [0, 0], [21, 22], [11, 12]] assert adapted["output_layer.linear.weight"][1] == [0, 0] assert adapted["output_layer.linear.weight"][-1] == [110, 111] assert adapted["output_layer.linear.bias"] == [200, 0, 202, 201] def test_masked_teacher_does_not_duplicate_an_already_final_mask() -> None: state = tensor_state(vocab=4) adapted = masked_to_absorbing( state, source_mask_index=3, target_vocab_size=4, expected_keys=set(BACKBONE_TENSOR_KEYS), ) assert adapted == state assert adapted is not state def test_teacher_adapters_allow_rotary_buffer_from_pinned_checkpoints() -> None: state = tensor_state() state["rotary_emb.inv_freq"] = [1, 2] adapted = masked_to_absorbing( state, source_mask_index=1, target_vocab_size=4, expected_keys=set(BACKBONE_TENSOR_KEYS), ) assert adapted["rotary_emb.inv_freq"] == [1, 2] def test_adapted_teacher_checkpoint_drops_lightning_metadata_for_weights_only_load( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: recipe = load_recipe("mdlm_di4c", "lm1b") launch = build_launch( recipe, root=tmp_path, source=tmp_path / recipe.source_path, output=tmp_path / "checkpoints" / "reference_reproduction" / "mdlm_di4c", teacher=tmp_path / "teacher.ckpt", devices=1, nodes=1, per_device_batch_size=1, seed=42, resume=True, ) launch.teacher.write_bytes(b"source checkpoint") state = tensor_state() for key in expected_backbone_keys() - set(state): state[key] = [0] source_payload = { "state_dict": state, "hyper_parameters": {"config": object()}, "callbacks": {"custom": object()}, } saved_payloads: list[dict[str, object]] = [] class FakeTorch: @staticmethod def load(path, **kwargs): return source_payload @staticmethod def save(payload, path): saved_payloads.append(payload) Path(path).write_bytes(b"adapted checkpoint") monkeypatch.setitem(sys.modules, "torch", FakeTorch) monkeypatch.setitem(recipes_module.BASE_VOCAB_SIZES, "lm1b", 3) recipes_module.adapt_teacher_checkpoint(launch, source_mask_index=1) assert len(saved_payloads) == 1 assert set(saved_payloads[0]) == {"state_dict"} assert "hyper_parameters" not in saved_payloads[0] @pytest.mark.parametrize( "mutation,match", [ (lambda state: state.pop("output_layer.linear.bias"), "keys"), (lambda state: state.update({"unexpected": [1]}), "keys"), ( lambda state: state.update({"output_layer.linear.weight": [[1], [2], [3]]}), "hidden", ), ], ) def test_teacher_adapters_fail_closed_on_key_or_shape_mismatch(mutation, match: str) -> None: state = tensor_state() mutation(state) with pytest.raises(RecipeError, match=match): uniform_to_absorbing(state, expected_keys=set(BACKBONE_TENSOR_KEYS)) @pytest.mark.parametrize( "wrapper", [ "scripts/train/flm.sh", "scripts/train/duo.sh", "scripts/train/mdlm.sh", "scripts/train/candi.sh", "scripts/train/rdlm.sh", "scripts/distill/fmlm.sh", "scripts/distill/duo_dcd.sh", "scripts/distill/mdlm_sdtt.sh", "scripts/distill/di4c.sh", ], ) def test_every_wrapper_supports_help(wrapper: str) -> None: completed = subprocess.run( ["bash", str(ROOT / wrapper), "--help"], cwd=ROOT, env={**os.environ, "DLB_PYTHON": sys.executable}, text=True, capture_output=True, check=False, ) assert completed.returncode == 0, completed.stderr assert "--dry-run" in completed.stdout def test_manifest_prerequisite_commands_dry_run_without_external_assets() -> None: commands = [ [ "bash", "scripts/train/candi.sh", "--source", "upstreams/candi", "--dataset", "owt", "--output", "checkpoints/reference_reproduction/candi/owt", "--dry-run", ], [ "bash", "scripts/distill/duo_dcd.sh", "--source", "upstreams/duo", "--dataset", "lm1b", "--teacher", "checkpoints/reference_reproduction/flm_baselines/lm1b/lm1b_Duo.ckpt", "--output", "checkpoints/reference_reproduction/duo_dcd/lm1b", "--rounds", "8", "--steps-per-round", "10000", "--global-batch-size", "128", "--learning-rate", "6e-5", "--dry-run", ], ] for command in commands: completed = subprocess.run( command, cwd=ROOT, env={**os.environ, "DLB_PYTHON": sys.executable}, text=True, capture_output=True, check=False, ) assert completed.returncode == 0, completed.stderr record = json.loads(completed.stdout) assert record["dry_run"] is True assert record["command"] assert Path(record["output"]).is_absolute() def test_fake_launch_publishes_before_run_and_resumes_exact_completion( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: recipe = load_recipe("flm", "lm1b") launch = build_launch( recipe, root=tmp_path, source=tmp_path / "upstreams" / "flm", output=tmp_path / "checkpoints" / "self_trained" / "lm1b" / "flm", teacher=None, devices=8, nodes=1, per_device_batch_size=64, seed=42, resume=True, ) monkeypatch.setattr(recipes_module, "_verify_source", lambda item: None) monkeypatch.setattr( recipes_module, "_prepare_training_cache", lambda item, root: { "mode": "fake", "processed_manifest_sha256": "a" * 64, }, ) monkeypatch.setattr( recipes_module, "_compose_config", lambda item, environment: "mode: train\n" ) calls = 0 def fake_run(command, **kwargs): nonlocal calls calls += 1 assert (launch.output / "recipe.json").is_file() assert (launch.output / "launch_argv.json").is_file() assert (launch.output / "provenance.json").is_file() checkpoint = launch.output / "checkpoints" / "last.ckpt" checkpoint.parent.mkdir(parents=True) checkpoint.write_bytes(b"fake checkpoint") return subprocess.CompletedProcess(command, 0) monkeypatch.setattr(recipes_module.subprocess, "run", fake_run) assert recipes_module.execute_launch(launch, root=tmp_path) == "completed" assert recipes_module.execute_launch(launch, root=tmp_path) == "skipped" assert calls == 1 completed = json.loads((launch.output / "completed.json").read_text()) assert completed["checkpoint"] == "model.ckpt" assert (launch.output / "model.ckpt").read_bytes() == b"fake checkpoint" def test_distilled_launch_canonicalizes_lightning_step_checkpoint( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: recipe = load_recipe("duo_di4c", "owt") teacher = tmp_path / "checkpoints" / "fixture" / "uniform_duo.ckpt" teacher.parent.mkdir(parents=True) teacher.write_bytes(b"fake teacher") launch = build_launch( recipe, root=tmp_path, source=tmp_path / "upstreams" / "di4c", output=tmp_path / "checkpoints" / "reference_reproduction" / "duo_di4c" / "owt", teacher=teacher, devices=2, nodes=1, per_device_batch_size=1, seed=42, resume=True, ) monkeypatch.setattr(recipes_module, "_verify_source", lambda item: None) monkeypatch.setattr( recipes_module, "_prepare_training_cache", lambda item, root: { "mode": "fake", "processed_manifest_sha256": "a" * 64, }, ) monkeypatch.setattr( recipes_module, "_compose_config", lambda item, environment: "mode: train\n" ) monkeypatch.setattr( recipes_module, "adapt_teacher_checkpoint", lambda item: {"mode": "fake_teacher_adapter"}, ) def fake_run(command, **kwargs): checkpoint = launch.output / "checkpoints" / "0-50000.ckpt" checkpoint.parent.mkdir(parents=True) checkpoint.write_bytes(b"fake distilled checkpoint") return subprocess.CompletedProcess(command, 0) monkeypatch.setattr(recipes_module.subprocess, "run", fake_run) assert recipes_module.execute_launch(launch, root=tmp_path) == "completed" completed = json.loads((launch.output / "completed.json").read_text()) assert completed["checkpoint"] == "student_checkpoints/50000.ckpt" assert (launch.output / completed["checkpoint"]).read_bytes() == ( b"fake distilled checkpoint" ) def test_training_cache_bridge_adds_attention_mask_for_duo_text_batches( tmp_path: Path, ) -> None: datasets = pytest.importorskip("datasets") Dataset = datasets.Dataset load_from_disk = datasets.load_from_disk processed = tmp_path / "data" / "processed" / "lm1b-bert-128" (tmp_path / "data" / "manifests").mkdir(parents=True) for split in ("train", "validation"): Dataset.from_dict( {"input_ids": [[101] + [7] * 126 + [102]]} ).save_to_disk(str(processed / split)) (tmp_path / "data" / "manifests" / "lm1b.json").write_text( json.dumps({"dataset": "lm1b", "sequence_length": 128}), encoding="utf-8", ) launch = build_launch( load_recipe("duo_dcd", "lm1b"), root=tmp_path, source=tmp_path / "upstreams" / "duo", output=tmp_path / "checkpoints" / "reference_reproduction" / "duo_dcd" / "lm1b", teacher=tmp_path / "checkpoints" / "fixture" / "uniform_duo.ckpt", devices=1, nodes=1, per_device_batch_size=128, seed=42, resume=True, ) record = recipes_module._prepare_training_cache(launch, tmp_path) cached = load_from_disk( str(launch.data_path / "lm1b_train_bs128_wrapped.dat") ) assert record["format"] == "hf_dataset_with_attention_mask_v1" assert set(cached.column_names) == {"input_ids", "attention_mask"} assert cached[0]["attention_mask"] == [1] * 128 def test_source_verification_ignores_python_bytecode_cache_status( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: recipe = load_recipe("mdlm_di4c", "lm1b") source = tmp_path / "upstreams" / "di4c" entrypoint = source / recipe.entrypoint entrypoint.parent.mkdir(parents=True) entrypoint.write_text("print('placeholder')\n", encoding="utf-8") launch = build_launch( recipe, root=tmp_path, source=source, output=tmp_path / "checkpoints" / "reference_reproduction" / "mdlm_di4c", teacher=tmp_path / "teacher.ckpt", devices=1, nodes=1, per_device_batch_size=1, seed=42, resume=True, ) def fake_run(command, **kwargs): if command[3] == "rev-parse": return subprocess.CompletedProcess( command, 0, stdout=recipe.source_commit + "\n", stderr="", ) if command[3] == "status": return subprocess.CompletedProcess( command, 0, stdout="?? sdtt/src/sdtt/__pycache__/\n?? sdtt/src/sdtt/core/__pycache__/\n", stderr="", ) raise AssertionError(command) monkeypatch.setattr(recipes_module.subprocess, "run", fake_run) recipes_module._verify_source(launch) def test_training_launch_disables_python_bytecode_writes( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: recipe = load_recipe("flm", "lm1b") launch = build_launch( recipe, root=tmp_path, source=tmp_path / "upstreams" / "flm", output=tmp_path / "checkpoints" / "self_trained" / "lm1b" / "flm", teacher=None, devices=8, nodes=1, per_device_batch_size=64, seed=42, resume=True, ) monkeypatch.setattr(recipes_module, "_verify_source", lambda item: None) monkeypatch.setattr( recipes_module, "_prepare_training_cache", lambda item, root: { "mode": "fake", "processed_manifest_sha256": "a" * 64, }, ) def fake_compose(item, environment): assert environment["PYTHONDONTWRITEBYTECODE"] == "1" return "mode: train\n" def fake_run(command, **kwargs): assert kwargs["env"]["PYTHONDONTWRITEBYTECODE"] == "1" checkpoint = launch.output / "checkpoints" / "last.ckpt" checkpoint.parent.mkdir(parents=True) checkpoint.write_bytes(b"fake checkpoint") return subprocess.CompletedProcess(command, 0) monkeypatch.setattr(recipes_module, "_compose_config", fake_compose) monkeypatch.setattr(recipes_module.subprocess, "run", fake_run) assert recipes_module.execute_launch(launch, root=tmp_path) == "completed"