| 58 | |
| 59 | |
| 60 | class Args(BaseSettings): |
| 61 | tgt_dir: Path | None = None |
| 62 | name: str | None = None |
| 63 | n_hidden_states: int | None = None |
| 64 | case: str = "all" |
| 65 | n_procs: int = 1 |
| 66 | disable_compile: bool = False |
| 67 | n_runs_warmup: int = 25 |
| 68 | n_runs_benchmark: int = 100 |
| 69 | n_samples: int = 1 |
| 70 | bench_fn: Literal["own", "nvbench", "fi-cupti"] = "fi-cupti" |
| 71 | nsys_profile: bool = False |
| 72 | top_k: int | None = None |
| 73 | top_p: float | None = None |
| 74 | |
| 75 | @model_validator(mode="after") |
| 76 | def _validate_distributed_bench_fn(self) -> "Args": |
| 77 | if self.n_procs > 1 and self.bench_fn == "nvbench": |
| 78 | raise ValueError( |
| 79 | "Distributed benchmarking is not supported with --bench_fn=nvbench. " |
| 80 | "nvbench controls iteration counts internally, which causes collective op " |
| 81 | "deadlocks when ranks run different numbers of iterations. " |
| 82 | "Use --bench_fn=own instead." |
| 83 | ) |
| 84 | return self |
| 85 | |
| 86 | def make_tp(self) -> TPInfo: |
| 87 | if self.n_procs > 1: |
| 88 | return TPInfo.from_world() |
| 89 | return TP1 |
| 90 | |
| 91 | def providers(self) -> list[str]: |
| 92 | if self.name is None or self.name == "default": |
| 93 | return DEFAULT_PROVIDERS |
| 94 | return self.name.split(",") |
| 95 | |
| 96 | |
| 97 | class CliArgs(Args, cli_parse_args=True): |
no outgoing calls