"""Parser2 DFlash2 双节点入口：验证实际命令和完整词表，不启动 GPU 进程。""" import argparse import json from pathlib import Path import numpy as np from torch.distributed.run import get_args_parser from speculators.train.cli import parse_vocab_mappings from speculators.train.config import TrainConfig from tests.unit.scripts.test_qwen3_6_two_node_launch import ( NETWORK_ENV, SCRIPT_DIR, assert_node_caches, assert_teacher_stopped, flag, read_capture, run_nodes, ) from tests.unit.scripts.test_qwen3_6_two_node_launch import ( launch_env as launch_env, # noqa: PLC0414 复用模拟启动 fixture。 ) SCRIPT = SCRIPT_DIR / "dflash2_infinity_parser2_flash_online_2node.sh" def test_parser2_two_node_launch(launch_env, monkeypatch): env = launch_env # 使用真实目录结构，确认传入的是 prepared Arrow 数据而不是 regen 根目录。 env.pop("DATA_DIR") data = ( Path(env["ROOT"]) / "datasets/infinity_parsers2_v2_1_max32768_vocab32k/dflash_data/full" ) data.mkdir(parents=True) for name in ("state.json", "dataset_info.json"): (data / name).write_text("{}") d2t = np.zeros(248320, dtype=np.int64) t2d = np.ones(248320, dtype=np.bool_) np.save(data / "d2t.npy", d2t) np.save(data / "t2d.npy", t2d) (Path(env["MODEL"]) / "config.json").write_text( json.dumps({"model_type": "qwen3_5", "text_config": {"vocab_size": 248320}}) ) env["NNODE_TEST_EXPECT_NODES"] = "2" for output, returncode in run_nodes("dflash2", env, (0, 1), script=SCRIPT): assert returncode == 0, output configs = [] for rank in (0, 1): teacher = read_capture(env, "vllm", rank) train = read_capture(env, "train", rank) assert teacher["env"]["CUDA_VISIBLE_DEVICES"] == "7,6" assert train["env"]["CUDA_VISIBLE_DEVICES"] == "5,4,3,2,1,0" assert "VLLM_MEDIA_LOADING_THREAD_COUNT" not in teacher["env"] assert "--api-server-count" not in teacher["argv"] for record in (teacher, train): for name in ("RANK", "WORLD_SIZE", "LOCAL_RANK", "LOCAL_WORLD_SIZE"): assert name not in record["env"] for name, value in NETWORK_ENV.items(): assert record["env"][name] == value assert "MASTER_ADDR" not in teacher["env"] assert "MASTER_PORT" not in teacher["env"] for name, value in { "--tensor-parallel-size": "1", "--data-parallel-size": "2", "--data-parallel-backend": "mp", "--nnodes": "1", "--node-rank": "0", "--master-addr": "127.0.0.1", "--data-parallel-address": "127.0.0.1", "--max-model-len": "65536", "--mm-processor-cache-gb": "0", "--served-model-name": env["MODEL"], }.items(): assert flag(teacher["argv"], name) == value monkeypatch.setenv("PET_NPROC_PER_NODE", "8") distributed = get_args_parser().parse_args(train["argv"]) assert distributed.nnodes == "2" assert distributed.nproc_per_node == "6" assert distributed.node_rank == rank assert distributed.master_addr == env["MASTER_ADDR"] assert str(distributed.master_port) == env["MASTER_PORT"] assert distributed.rdzv_backend == "static" assert not distributed.standalone assert not distributed.no_python cfg = TrainConfig.resolve(distributed.training_script_args).flatten() configs.append(cfg) assert cfg["run_name"] == "dflash2-parser2_1-2node" run_dir = ( Path(env["ROOT"]) / "model_weights/dflash2_parser2_1_flash_2node" / cfg["run_name"] ) assert cfg["save_path"] == str(run_dir / "checkpoints") assert cfg["log_dir"] == str(run_dir) assert_node_caches(teacher, train, env, rank, cfg["run_name"]) for name, value in { "speculator_type": "dflash2", "checkpoint_freq": 0.1, "verifier_name_or_path": env["MODEL"], "data_path": str(data), "draft_vocab_size": None, "draft_arch": "qwen3", "num_layers": 5, "mask_token_id": 248077, "target_layer_ids": [2, 7, 12, 17, 22], "draft_mrope_full_head_hack": True, "sliding_window": 2048, "sliding_window_non_causal": True, "full_attention_indices": [], "total_seq_len": 16384, "block_size": 16, "max_anchors": 1024, "sample_from_anchor": None, "loss_fn": "ce", "per_position_loss_weight": "dpace", "conv_kernel_size": 2, "conv_group_size": 16, "selector_rank": 256, "selector_top_k": 16, "selector_loss_alpha": 0.1, "num_workers": 12, "prefetch_factor": 4, "fetch_threads": 1, "dataloader_in_order": True, "vllm_http_keepalive": True, "request_timeout": 120, "max_retries": 3, "generation_validation_retries": 2, "max_consecutive_generation_failures": 20, "fail_on_hidden_state_error": False, }.items(): assert cfg[name] == value, name assert cfg["vllm_endpoint"] == f"http://127.0.0.1:{env['VLLM_PORT']}/v1" loaded_d2t, loaded_t2d, vocab_size = parse_vocab_mappings( argparse.Namespace(**cfg) ) assert vocab_size == 248320 np.testing.assert_array_equal(loaded_d2t.numpy(), d2t) np.testing.assert_array_equal(loaded_t2d.numpy(), t2d) hs_path = Path(cfg["hidden_states_path"]) assert hs_path.parent == Path("/tmp") assert not hs_path.exists() # 脚本退出时清理本次目录。 assert str(hs_path) == flag(teacher["argv"], "--hidden-states-path") assert_teacher_stopped(teacher) for key in ("save_path", "run_name", "log_dir"): assert configs[0][key] == configs[1][key] assert configs[0]["hidden_states_path"] != configs[1]["hidden_states_path"] # 训练直接加载全词表映射，不改写数据目录。 np.testing.assert_array_equal(np.load(data / "d2t.npy"), d2t) np.testing.assert_array_equal(np.load(data / "t2d.npy"), t2d) def test_missing_parser2_prepared_data_fails_before_teacher(launch_env): launch_env.pop("DATA_DIR") [(output, returncode)] = run_nodes("dflash2", launch_env, (0,), script=SCRIPT) assert returncode != 0 assert "dflash_data/full/state.json" in output assert not list(Path(launch_env["NNODE_TEST_CAPTURE"]).iterdir())