| """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, |
| ) |
|
|
| SCRIPT = SCRIPT_DIR / "dflash2_infinity_parser2_flash_online_2node.sh" |
|
|
|
|
| def test_parser2_two_node_launch(launch_env, monkeypatch): |
| env = launch_env |
| |
| 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()) |
|
|
