From 31f2e839c53c5a1622fae7eaf49562f8d5591e63 Mon Sep 17 00:00:00 2001 From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> Date: Thu, 6 Aug 2026 23:39:20 -0700 Subject: [PATCH 1/2] [None][infra] Add execution and test-runner skills for Claude Code Restructure the .claude toolkit around a layered execution model and add a test-runner stack, so skills that need to run something delegate to a single executor instead of each embedding its own Slurm/Docker plumbing. Execution layer: - exec-env-check: probe GPUs, Docker, and Slurm access, return the execution scenario the callers branch on - exec-local-docker, exec-local-slurm, exec-remote-slurm: workflow-agnostic executors driven by a job_spec.json contract, with matching agents for the two Slurm variants - exec-slurm-compile: rename mount_dir to user_root_dir so the bind-mount parameter is not confused with container mount syntax Test-runner stack: - trtllm-test-specialist: classify module vs model tests, build commands, and delegate all execution to trtllm-case-executor - trtllm-case-executor and trtllm-test-script-builder: environment selection, script generation, and job submission - trtllm-test-specialist agent as the entry point Performance: - perf-optimization-casebook: 40 past optimizations recorded as decision precedents with machine-readable frontmatter, tag and pattern registries - perf-analysis and perf-optimization consult the casebook for prior art before recommending or routing an optimization Consolidation: - fold ad-sharding-ir-port into ad-model-onboard, which now covers the full sharding-aware IR porting procedure - drop trtllm-model-onboard-multimodal and the imported-kernel ABI checklist from trtllm-moe-develop Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com> --- .claude/agents/exec-local-slurm.md | 31 + .claude/agents/exec-remote-slurm.md | 40 + .claude/agents/trtllm-test-specialist.md | 18 + .claude/skills/ad-model-onboard/SKILL.md | 137 +- .claude/skills/ad-sharding-ir-port/SKILL.md | 251 ---- .claude/skills/exec-env-check/SKILL.md | 142 ++ .claude/skills/exec-local-docker/SKILL.md | 121 ++ .claude/skills/exec-local-slurm/SKILL.md | 443 ++++++ .claude/skills/exec-remote-slurm/SKILL.md | 970 ++++++++++++++ .../skills/exec-remote-slurm/scripts/build.sh | 42 + .../exec-remote-slurm/scripts/build.slurm | 41 + .claude/skills/exec-slurm-compile/SKILL.md | 14 +- .../exec-slurm-compile/scripts/compile.slurm | 6 +- .../scripts/submit_compile.sh | 6 +- .claude/skills/perf-analysis/SKILL.md | 6 + .../perf-optimization-casebook/SKILL.md | 293 ++++ .../data/aliases.yaml | 53 + .../data/patterns.yaml | 406 ++++++ .../perf-optimization-casebook/data/tags.yaml | 116 ++ .../references/case-template.md | 211 +++ .../references/communication/deepep.md | 37 + .../references/communication/index.md | 74 ++ .../communication/low-precision-dispatch.md | 38 + .../communication/mnnvl-twoshot-allreduce.md | 35 + .../shape-aware-allreduce-autotune.md | 33 + .../userbuffers-symmetric-memory.md | 37 + .../fold-scale-swizzle-into-kernel.md | 33 + .../kernel-and-fusion/fp8-mla-kv-cache.md | 38 + .../kernel-and-fusion/fuse-add-norm-quant.md | 39 + .../kernel-and-fusion/fuse-ar-epilogue.md | 38 + .../fuse-datamovement-into-quantize.md | 85 ++ .../fuse-moe-routing-kernel.md | 115 ++ .../kernel-and-fusion/fuse-qk-norm-rope.md | 36 + .../hw-matched-lowprec-moe-gemm.md | 36 + .../references/kernel-and-fusion/index.md | 158 +++ .../mega-fuse-moe-deepgemm.md | 41 + .../ranking-only-precision-tf32.md | 87 ++ .../reevaluate-fusion-boundary-per-dtype.md | 87 ++ .../kernel-and-fusion/relax-tl-constexpr.md | 69 + .../sparse-mla-topk-attention.md | 117 ++ .../specialize-topk-selection-kernel.md | 103 ++ .../split-mla-reduction-kernel.md | 35 + .../kernel-and-fusion/triton-to-cpp-op.md | 81 ++ .../trtllm-gen-fp4-moe-backend.md | 36 + .../runtime-execution/attention-dp-padding.md | 33 + .../auxiliary-cache-in-kv-manager.md | 94 ++ .../cache-step-invariant-per-layer.md | 85 ++ .../chunked-prefill-aligned-auxiliary.md | 86 ++ .../runtime-execution/cuda-graph-padding.md | 33 + .../free-mla-intermediates.md | 33 + .../hoist-torch-compile-closures.md | 77 ++ .../references/runtime-execution/index.md | 165 +++ .../runtime-execution/mla-kv-cache-reuse.md | 35 + .../move-bookkeeping-into-cpp-op.md | 36 + .../multi-stream-shared-routed-expert.md | 37 + .../overlap-mla-rope-uk-bgemm.md | 35 + .../runtime-execution/overlap-online-eplb.md | 35 + .../runtime-execution/overlap-scheduler.md | 41 + .../references/runtime-execution/pdl.md | 35 + .../runtime-execution/piecewise-cuda-graph.md | 35 + .../pybind-wrapper-pure-python.md | 93 ++ .../relaxed-mtp-acceptance.md | 37 + .../skip-sparse-path-when-degenerate.md | 90 ++ .../split-custom-op-for-piecewise-capture.md | 93 ++ .../runtime-execution/two-model-mtp-eagle.md | 35 + .claude/skills/perf-optimization/SKILL.md | 61 +- .claude/skills/trtllm-case-executor/SKILL.md | 440 ++++++ .../scripts/detect_slurm_env.sh | 449 +++++++ .../trtllm-model-onboard-multimodal/SKILL.md | 471 ------- .claude/skills/trtllm-moe-develop/SKILL.md | 20 - .../trtllm-test-script-builder/SKILL.md | 385 ++++++ .../references/trtllm_test_template.md | 520 ++++++++ .../scripts/build_slurm_script.py | 643 +++++++++ .../scripts/slurm_run_custom.sh | 33 + .../skills/trtllm-test-specialist/SKILL.md | 681 ++++++++++ .../references/agg_config_template.yaml | 257 ++++ .../references/benchmark_config_template.yaml | 133 ++ .../references/disagg_config_template.yaml | 63 + .../references/smoke_test_config_template.yml | 19 + .../references/test_config_template.yaml | 227 ++++ .../trtllm_test_fix_recommendations.md | 14 + .../scripts/build_test_command.py | 768 +++++++++++ .../scripts/extract_test_markers.py | 200 +++ .../scripts/generate_benchmark_config.py | 1184 +++++++++++++++++ .../scripts/generate_report.py | 297 +++++ .../scripts/parse_config.py | 237 ++++ 86 files changed, 12373 insertions(+), 807 deletions(-) create mode 100644 .claude/agents/exec-local-slurm.md create mode 100644 .claude/agents/exec-remote-slurm.md create mode 100644 .claude/agents/trtllm-test-specialist.md delete mode 100644 .claude/skills/ad-sharding-ir-port/SKILL.md create mode 100644 .claude/skills/exec-env-check/SKILL.md create mode 100644 .claude/skills/exec-local-docker/SKILL.md create mode 100644 .claude/skills/exec-local-slurm/SKILL.md create mode 100644 .claude/skills/exec-remote-slurm/SKILL.md create mode 100644 .claude/skills/exec-remote-slurm/scripts/build.sh create mode 100644 .claude/skills/exec-remote-slurm/scripts/build.slurm create mode 100644 .claude/skills/perf-optimization-casebook/SKILL.md create mode 100644 .claude/skills/perf-optimization-casebook/data/aliases.yaml create mode 100644 .claude/skills/perf-optimization-casebook/data/patterns.yaml create mode 100644 .claude/skills/perf-optimization-casebook/data/tags.yaml create mode 100644 .claude/skills/perf-optimization-casebook/references/case-template.md create mode 100644 .claude/skills/perf-optimization-casebook/references/communication/deepep.md create mode 100644 .claude/skills/perf-optimization-casebook/references/communication/index.md create mode 100644 .claude/skills/perf-optimization-casebook/references/communication/low-precision-dispatch.md create mode 100644 .claude/skills/perf-optimization-casebook/references/communication/mnnvl-twoshot-allreduce.md create mode 100644 .claude/skills/perf-optimization-casebook/references/communication/shape-aware-allreduce-autotune.md create mode 100644 .claude/skills/perf-optimization-casebook/references/communication/userbuffers-symmetric-memory.md create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fold-scale-swizzle-into-kernel.md create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fp8-mla-kv-cache.md create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-add-norm-quant.md create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-ar-epilogue.md create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-datamovement-into-quantize.md create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-moe-routing-kernel.md create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/fuse-qk-norm-rope.md create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/hw-matched-lowprec-moe-gemm.md create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/index.md create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/mega-fuse-moe-deepgemm.md create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/ranking-only-precision-tf32.md create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/reevaluate-fusion-boundary-per-dtype.md create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/relax-tl-constexpr.md create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/sparse-mla-topk-attention.md create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/specialize-topk-selection-kernel.md create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/split-mla-reduction-kernel.md create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/triton-to-cpp-op.md create mode 100644 .claude/skills/perf-optimization-casebook/references/kernel-and-fusion/trtllm-gen-fp4-moe-backend.md create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/attention-dp-padding.md create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/auxiliary-cache-in-kv-manager.md create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/cache-step-invariant-per-layer.md create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/chunked-prefill-aligned-auxiliary.md create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/cuda-graph-padding.md create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/free-mla-intermediates.md create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/hoist-torch-compile-closures.md create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/index.md create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/mla-kv-cache-reuse.md create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/move-bookkeeping-into-cpp-op.md create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/multi-stream-shared-routed-expert.md create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/overlap-mla-rope-uk-bgemm.md create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/overlap-online-eplb.md create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/overlap-scheduler.md create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/pdl.md create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/piecewise-cuda-graph.md create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/pybind-wrapper-pure-python.md create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/relaxed-mtp-acceptance.md create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/skip-sparse-path-when-degenerate.md create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/split-custom-op-for-piecewise-capture.md create mode 100644 .claude/skills/perf-optimization-casebook/references/runtime-execution/two-model-mtp-eagle.md create mode 100644 .claude/skills/trtllm-case-executor/SKILL.md create mode 100755 .claude/skills/trtllm-case-executor/scripts/detect_slurm_env.sh delete mode 100644 .claude/skills/trtllm-model-onboard-multimodal/SKILL.md create mode 100644 .claude/skills/trtllm-test-script-builder/SKILL.md create mode 100644 .claude/skills/trtllm-test-script-builder/references/trtllm_test_template.md create mode 100755 .claude/skills/trtllm-test-script-builder/scripts/build_slurm_script.py create mode 100755 .claude/skills/trtllm-test-script-builder/scripts/slurm_run_custom.sh create mode 100644 .claude/skills/trtllm-test-specialist/SKILL.md create mode 100644 .claude/skills/trtllm-test-specialist/references/agg_config_template.yaml create mode 100644 .claude/skills/trtllm-test-specialist/references/benchmark_config_template.yaml create mode 100644 .claude/skills/trtllm-test-specialist/references/disagg_config_template.yaml create mode 100644 .claude/skills/trtllm-test-specialist/references/smoke_test_config_template.yml create mode 100644 .claude/skills/trtllm-test-specialist/references/test_config_template.yaml create mode 100644 .claude/skills/trtllm-test-specialist/references/trtllm_test_fix_recommendations.md create mode 100644 .claude/skills/trtllm-test-specialist/scripts/build_test_command.py create mode 100644 .claude/skills/trtllm-test-specialist/scripts/extract_test_markers.py create mode 100644 .claude/skills/trtllm-test-specialist/scripts/generate_benchmark_config.py create mode 100755 .claude/skills/trtllm-test-specialist/scripts/generate_report.py create mode 100644 .claude/skills/trtllm-test-specialist/scripts/parse_config.py diff --git a/.claude/agents/exec-local-slurm.md b/.claude/agents/exec-local-slurm.md new file mode 100644 index 000000000000..795f5751441c --- /dev/null +++ b/.claude/agents/exec-local-slurm.md @@ -0,0 +1,31 @@ +--- +name: exec-local-slurm +description: > + Execute a TensorRT-LLM workload on a local Slurm cluster. Supports persistent + allocation (allocate once via nohup salloc, reuse across runs) and one-shot + sbatch. Workflow-agnostic — handles pytest, eval, benchmark, and custom + scripts identically. The orchestrator (typically trtllm-case-executor) writes + a job spec to /job_spec.json and invokes this agent to run it. +tools: ["Read", "Write", "Edit", "Bash", "Grep", "Glob"] +model: sonnet +license: Apache-2.0 +--- + +You are the local Slurm executor agent. Load the `exec-local-slurm` skill (`trtllm-agent-toolkit:exec-local-slurm`) and follow its procedure exactly. The skill is the single source of truth for the execution flow; this agent file only adds the contract with the caller and the invariants that must hold across every run. + +## Input + +The caller passes a path to a job spec — typically `/job_spec.json` — plus a short summary of the fields used in this run. **Read `job_spec.json` first.** The skill's "Input (from orchestrator prompt)" section enumerates every field it consumes (`script_path`, `work_dir`, `model_name`, `workflow_type`, `success_patterns`, `failure_patterns`, `log_file_pattern`, `monitor_timeout_seconds`, `persistent_mode`, `release_allocation`, `alloc_time_limit`, `docker_image`, `container_name`, `container_mounts`, `repo_root`, `slurm_params`). + +Do not re-derive any field that `trtllm-case-executor` already wrote into the spec. + +## Invariants + +- **Hang detection.** Poll the log periodically; on a case-insensitive `hang detected` match, kill the process group and report `HANG_DETECTED`. Implementation lives in the skill. +- **Wall-clock limit.** Honor `monitor_timeout_seconds` (default `3600`). On timeout, kill and report `TIMEOUT`. +- **Persistent allocation lifecycle.** Default `persistent_mode=true`; reuse an existing allocation when present and valid. Only release when `release_allocation=true` is explicitly set — never auto-release. +- **Single source of truth.** `node_count`, `job_name`, `container_image`, and `slurm_params` come from `job_spec.json`. Never recompute them or re-grep `current_image_tags.properties`. + +## Output + +Return a single report to the caller with: task type, status (`PASSED` / `FAILED` / `TIMEOUT` / `HANG_DETECTED` / `OUT_OF_MEMORY` / `CANCELLED` / `ERROR` / `BUILD_FAILED`), Slurm job id (and allocation id when persistent), log file path, summary, and any error excerpts (last ~100 lines on build/job failure). Do not perform follow-up actions beyond what the skill prescribes. diff --git a/.claude/agents/exec-remote-slurm.md b/.claude/agents/exec-remote-slurm.md new file mode 100644 index 000000000000..496493d0aa56 --- /dev/null +++ b/.claude/agents/exec-remote-slurm.md @@ -0,0 +1,40 @@ +--- +name: exec-remote-slurm +description: > + Execute a TensorRT-LLM workload on a remote Slurm cluster via SSH. Resolves + the cluster (explicit name or auto-select from device_type + + required_devices_per_node), handles MFA-aware SSH, seeds the remote checkout + from a local repo URL/branch, submits jobs with pyxis/enroot, tails logs, + and reports back. The orchestrator (typically trtllm-case-executor) writes a + job spec to /job_spec.json and invokes this agent to run it. +tools: ["Read", "Write", "Edit", "Bash", "Grep", "Glob"] +model: sonnet +license: Apache-2.0 +--- + +You are the remote Slurm executor agent. Load the `exec-remote-slurm` skill (`trtllm-agent-toolkit:exec-remote-slurm`) and follow its procedure exactly. The skill is the single source of truth for the execution flow; this agent file only adds the contract with the caller and the invariants that must hold across every run. + +## Input + +The caller passes a path to a job spec — typically `/job_spec.json` — plus a short summary of the fields used in this run. **Read `job_spec.json` first.** The skill's "Input" section enumerates every field it consumes (`script_path`, `script_name`, `work_dir`, `model_name`, `workflow_type`, `success_patterns`, `failure_patterns`, `log_file_pattern`, `slurm_cluster`, `ssh_host`, `slurm_user`, `remote_cwd`, `remote_work_dir`, `slurm_password`, `extra_files`, `repo_url`, `repo_branch`, `device_type`, `total_required_devices`, `required_devices_per_node`, `container_image`, `node_count`, `job_name`, `monitor_timeout_seconds`). + +Do not re-derive any field that `trtllm-case-executor` already wrote into the spec. + +## Cluster resolution + +**Precondition — optional dependency.** Before either bullet below, check whether the `skills/internal-env-info/` directory exists in the toolkit. If it does **not**, do not attempt to load any reference file or invoke the skill. Proceed using only the cluster fields the orchestrator wrote into `job_spec.json` (`ssh_host`, `slurm_user`, `remote_cwd`, `partition`, `account`, `container_image`, `mounts`, `gpus_per_node`, `mfa_style`). If a required field is also absent, stop and ask the user to supply it — do **not** report the missing skill as an error. + +- **Explicit cluster** (`slurm_cluster` is set) → invoke the `internal-env-info` skill in single-cluster mode to fetch per-cluster info (`mfa_style`, `default_models_repo`, `default_user_root_dir`, `gpus_per_node`). Connection fields (`ssh_host`, `slurm_user`, `remote_cwd`, `account`, `partition`, `mounts`, `container_image`) come from `job_spec.json` — there is no per-cluster connection-config file to parse. +- **Auto-select** (only hardware constraints are present) → invoke the `internal-env-info` skill in constraint-based mode, passing `device_type` and `required_devices_per_node` (and optionally `total_required_devices`). Never re-implement the constraint filter. + +## Invariants + +- **`mfa_style` decides the SSH path.** `false` → direct; `true` → MFA flow; `null` → probe direct first, then fall back, and ask the user to update the cluster reference. Do **not** probe-and-fall-back on the SSH error string when `mfa_style` is known. +- **Hang detection + `monitor_timeout_seconds`.** Same semantics as local execution; implementation lives in the skill. +- **Pass-through fields.** `container_image`, `node_count`, and `job_name` come from `job_spec.json` and are used verbatim. Apply transport-specific URL rewrites (e.g., enroot `/` → `#`) at use-time. Never re-grep `current_image_tags.properties` or reconstruct `job_name` from `account` / `subproject` / `detail`. +- **Remote repo bootstrap.** Use `repo_url` and `repo_branch` from `job_spec.json` to ensure the remote checkout matches the local one before submission. +- **`node_count` is authoritative.** Use it directly as `--nodes`; never recompute from totals. + +## Output + +Return a single report to the caller with: task type, status (`PASSED` / `FAILED` / `TIMEOUT` / `HANG_DETECTED` / `OUT_OF_MEMORY` / `CANCELLED` / `ERROR` / `BUILD_FAILED`), remote Slurm job id, remote log path (and a local copy when synced back), summary, and any error excerpts (last ~100 lines on build/job failure). Do not perform follow-up actions beyond what the skill prescribes. diff --git a/.claude/agents/trtllm-test-specialist.md b/.claude/agents/trtllm-test-specialist.md new file mode 100644 index 000000000000..0952a24f01c2 --- /dev/null +++ b/.claude/agents/trtllm-test-specialist.md @@ -0,0 +1,18 @@ +--- +name: trtllm-test-specialist +description: > + Runs model-level and module-level tests for TensorRT-LLM. Classifies the test + scope (module test or model test), builds the appropriate test commands, and + delegates execution to trtllm-case-executor. Supports functionality/smoke + tests, benchmarks, and evaluations. Writes structured test reports to a + caller-specified path. +tools: ["Read", "Write", "Edit", "Bash", "Grep", "Glob"] +license: Apache-2.0 +--- + +Role: dispatch TRT-LLM model-level and module-level test requests. + +Load the `trtllm-agent-toolkit:trtllm-test-specialist` skill, pass the caller's parameters through verbatim, and return its result. + +- If the caller supplies `report_file`, write the report to that exact path — do not substitute the skill's default (`./-auto-test-report.md`) or invent your own. +- Return the skill's status and final report path verbatim; do not re-summarize or rename. diff --git a/.claude/skills/ad-model-onboard/SKILL.md b/.claude/skills/ad-model-onboard/SKILL.md index 2e7de8db5a65..e48c56875677 100644 --- a/.claude/skills/ad-model-onboard/SKILL.md +++ b/.claude/skills/ad-model-onboard/SKILL.md @@ -340,9 +340,142 @@ GH_CONFIG_DIR= gh pr view --json reviews,state **Do NOT stop polling prematurely.** The loop must continue until the PR is approved or a clear termination signal is received. If polling has been running for an extended period (e.g., >2 hours) with no new activity, inform the user that you are still monitoring and ask if they want you to continue or stop. -## Sharding-aware IR model porting +## Sharding-aware IR model porting (`modeling_*_ir.py`) -For porting an existing custom model to a sharding-aware `_ir.py` variant, see the `ad-sharding-ir-port` skill. +Use this when porting an existing AutoDeploy custom model (`tensorrt_llm/_torch/auto_deploy/models/custom/modeling_*.py`) to explicit sharding hint ops in `modeling_*_ir.py` **in the same directory** (no separate `new_sharding/` tree). The exported FX graph must fully specify how the model should be sharded: the `apply_sharding_hints` transform combines hints with a runtime `DistConfig` for deterministic, node-local sharding. + +**Argument reference:** Do not duplicate operator tables here. Refer to the custom op docstrings in `tensorrt_llm/_torch/auto_deploy/custom_ops/` for the complete argument reference (including sharding hints, `tp_mode`, `layer_type`, and which ops accept hints). + +### Reference examples (study before porting) + +| Original | IR / sharding-aware | Layer types | +|----------|---------------------|-------------| +| `modeling_nemotron_h.py` | `modeling_nemotron_h_ir.py` | Mamba SSM, MHA, SwiGLU MLP, MoE | +| `modeling_qwen3_5_moe.py` | `modeling_qwen3_5_moe_ir.py` | GatedDeltaNet, Gated MHA, SwiGLU MLP, MoE | +| `modeling_mistral.py` | `modeling_mistral_ir.py` | MHA, SwiGLU MLP (simplest) | +| `modeling_deepseek_v2.py` | `modeling_deepseek_v2_ir.py` | MLA, SwiGLU MLP, MoE | + +### Step-by-step porting procedure + +#### Step 1: Copy the source file + +```bash +cp tensorrt_llm/_torch/auto_deploy/models/custom/modeling_foo.py \ + tensorrt_llm/_torch/auto_deploy/models/custom/modeling_foo_ir.py +``` + +#### Step 2: Update the module docstring and add imports + +At the top of the IR file: + +```python +import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 -- register all ops +``` + +Do **not** add global `SHARD_*` flags. Layer-level control uses the `layer_type` hint on each op and `shard_layers` in YAML. + +#### Step 3: Replace linear projections + +For every `self.proj(x)` or `nn.Linear` call, use `torch.ops.auto_deploy.torch_linear_simple` with explicit `tp_mode` and `layer_type`. Always set `tp_mode` unconditionally (no `if _s else "none"`). **Rules:** opening projections (Q/K/V/gate/up/in_proj) → `"colwise"`; closing (O/down/out_proj) → `"rowwise"`; tiny outputs (e.g. `shared_expert_gate` dim 1) → `"none"`; MLA latent projections (q_a, kv_a) → `"none"`. For fused weights split later, pass `output_sizes=[...]`. For GQA, use `tp_min_local_shape=self.head_dim` on K/V colwise lines. + +#### Step 4: Replace split / chunk after fused colwise projections + +Use `torch.ops.auto_deploy.split_with_sizes` with `shardable` / `layer_type` where sizes scale with TP. + +#### Step 5: Replace view / reshape with concrete head counts + +During `torch.export`, `-1` becomes concrete; after TP, wrong values break. Any reshape whose dimension is a head count that scales with TP must use `torch.ops.auto_deploy.view` with `tp_scaled_dim` set appropriately. Safe cases: flat-to-2D, or `[B,S,-1]` when the input is already correctly sharded. + +#### Step 6: Insert `all_reduce` + +After every rowwise projection, add `torch.ops.auto_deploy.all_reduce(..., layer_type=...)`. **Parallel branch rule:** when branches merge by addition, use a **single** `all_reduce` after the sum (e.g. MoE routed + shared expert; parallel attention + MLP residual branches). + +#### Step 7: Special ops (Conv1d, SSM, GatedDeltaNet, gated RMSNorm) + +Add sharding hints on `torch_causal_conv1d`, `torch_ssm`, `torch_gated_delta_rule`, `torch_rmsnorm_gated` per docstrings—typically `shardable` / `output_sizes` / `tp_mode` as required. + +#### Step 8: MoE + +Pass `layer_type="moe"` into `torch_moe`; `apply_sharding_hints` handles EP/TP. + +#### Step 9: Register the IR model + +1. Bottom of the IR file: `AutoModelForCausalLMFactory.register_custom_model_cls("ConfigClassName", ForCausalLM)` (same pattern as Phase 4). +2. Add a **side-effect import** in `tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py` (e.g. `from . import modeling_foo_ir # noqa: F401`) and extend `__all__` if you export symbols. Without this import, worker processes may not load your class and `apply_sharding_hints` can report **0 nodes processed**. Do **not** use a separate `register_sharded_models.py` indirection. + +#### Step 10: YAML — composable registry pattern + +Prefer the model registry (`examples/auto_deploy/model_registry/models.yaml`) and **compose** shared fragments under `examples/auto_deploy/model_registry/configs/`, same as other models: list `dashboard_default.yaml`, the right `world_size_N.yaml`, then a dedicated fragment (e.g. `enable_sharder_ir.yaml`) that holds IR sharding transforms. That fragment should disable legacy sharding passes and enable hint-driven sharding. Registry fragments are deep-merged in `yaml_extra` order (see `DynamicYamlMixInForSettings` in `tensorrt_llm/_torch/auto_deploy/utils/_config.py`); place transform keys under `transforms:` so they merge with `dashboard_default.yaml`. Standalone experiment YAMLs for `build_and_run_ad` may wrap the same fields under a top-level `args:` block matching `LlmArgs`. + +Example transform block: + +```yaml +# Typical contents for enable_sharder_ir.yaml (registry composable fragment) +transforms: + export_to_gm: + num_moe_experts_for_export: 2 # often required when expert count is large (>64) + detect_sharding: + stage: sharding + enabled: false + sharding_transform_executor: + stage: sharding + enabled: false + apply_sharding_hints: + stage: sharding + enabled: true + run_shape_prop: true + allreduce_strategy: NCCL + # shard_layers: ['mha', 'mlp'] # optional selective sharding + gather_logits_before_lm_head: + enabled: true +``` + +Use `world_size: 8` when validating TP head-divisibility. Optional `shard_layers` limits which `layer_type` hints are processed; unset means shard all shardable nodes. + +#### Step 11: Validate + +Do not report success until a run completes successfully. + +1. Prefer `python examples/auto_deploy/build_and_run_ad.py --model --use-registry` after adding/updating the registry entry and composable YAMLs (Phase 8–9 style). +2. `apply_sharding_hints` logs should show **`N nodes processed` with N > 0**. +3. If validation fails with infrastructure limits (e.g. head count not divisible by `world_size`), document the assert and compatible sizes; do not "fix" core `sharding.py` / custom op schemas without owner review. +4. If blocked by missing infrastructure support, rename artifacts to `broken_modeling_*_ir.py` / broken YAML and file a short error report for humans (do not silently patch core transforms). + +**Layer type strings** (for `layer_type` / `shard_layers`): use `"mha"`, `"mla"`, `"mlp"`, `"moe"`, `"ssm"`, `"delta"`, or `"unknown"` (default; skipped when `shard_layers` is set). Match the conventions used in `apply_sharding_hints` and project enums. + +### Layer-specific sharding patterns + +**MHA (standard or gated):** `layer_type="mha"`: q/k/v colwise (GQA: `tp_min_local_shape`), `view` with `tp_scaled_dim` for head dim, o rowwise + `all_reduce`. Fused Q+gate interleaved per head: colwise without `output_sizes`; contiguous Q|K|V fused blocks need `output_sizes`. + +**SwiGLU MLP:** `layer_type="mlp"`: gate/up colwise, down rowwise + `all_reduce`. + +**Mamba / SSM:** `layer_type="ssm"`: in_proj colwise + `output_sizes`, splits shardable, conv1d shardable + `output_sizes`, views, `torch_ssm` shardable, norm gated colwise if weight scales, out rowwise + `all_reduce`. + +**GatedDeltaNet:** `layer_type="delta"`: in_proj_qkv with `output_sizes`, other in_projs colwise, conv1d/splits/views as above, `torch_gated_delta_rule` shardable, out rowwise + `all_reduce`. + +**MoE + shared expert:** `layer_type="moe"`: router replicated; one `all_reduce` after `routed + shared`, not two. + +**MLA (DeepSeek):** `layer_type="mla"`: keep `torch_mla` intact with `shardable=True`—do **not** decompose into separate linears + `torch_attention` (introduces bad `expand`/`view` with concrete head counts). q_a/kv_a latent: `tp_mode="none"`; q_b colwise; `o_proj` rowwise + `all_reduce`. + +### Common pitfalls (sharding IR) + +1. **Missing `auto_deploy::view` for head reshapes** — concrete shapes from export break after sharding. +2. **Sharding tiny projections** — dim-1 gates: `tp_mode="none"`. +3. **Double `all_reduce` in MoE** — one merge-point reduction for routed + shared. +4. **Cross-layer parameter contamination** — in `_apply_hint_*` handlers using `get_source_nodes()`, restrict with `allowed_ops` so residual links do not pull weights from other layers. +5. **Missing `num_moe_experts_for_export`** for very large expert counts — export can hang. +6. **Decomposing ops that absorb weights** (e.g. `torch_mla`) — use `shardable` + handler instead of splitting into plain linears. +7. **Interleaved vs contiguous fused weights** — interleaved per-head groups: colwise only; contiguous Q|K|V blocks: require `output_sizes`. +8. **Omitting `layer_type` when using `shard_layers`** — `"unknown"` nodes are skipped; set hints explicitly on sharding-aware ops. +9. **`layer_type` on non-hint ops** — do **not** pass `layer_type` to ops that are not designed for sharding hints (e.g. `torch_attention`, `torch_l2norm`, `torch_rope_*`); extra positional args break calls. Confirm in `custom_ops/` docstrings which ops accept hints. +10. **Conditional hint values** — no `if _s else "none"`; use unconditional hints and rely on `shard_layers` / transform config. + +### Sharding IR validation checklist (human review) + +- `world_size=1`: unsharded path; hints should not break correctness. +- `world_size=2` and `8`: shape checks and coherent output. +- `apply_sharding_hints` node count vs expectation. +- Optional: `shard_layers: ['moe']` to verify selective sharding. ## Key Gotchas - **Canonical ops first:** Always use `torch.ops.auto_deploy.torch_*` canonical ops whenever one exists for the operation. This is how AD knows what to optimize. Writing manual attention, MoE, RoPE, or normalization in plain PyTorch instead of using the canonical op will prevent AD transforms from working. diff --git a/.claude/skills/ad-sharding-ir-port/SKILL.md b/.claude/skills/ad-sharding-ir-port/SKILL.md deleted file mode 100644 index 1eb40bdbb380..000000000000 --- a/.claude/skills/ad-sharding-ir-port/SKILL.md +++ /dev/null @@ -1,251 +0,0 @@ ---- -name: ad-sharding-ir-port -description: > - Adds sharding-aware IR hints (op substitutions, sharding kwargs, all_reduce - insertions) directly into an existing AutoDeploy custom model - (modeling_*.py). Edits the file in place — no separate _ir.py copy. - Validates with apply_sharding_hints and end-to-end multi-GPU runs. -license: Apache-2.0 -metadata: - author: NVIDIA Corporation ---- - -# Adding Sharding IR Hints to an AutoDeploy Custom Model - -**Input:** An existing AutoDeploy custom model at `tensorrt_llm/_torch/auto_deploy/models/custom/modeling_*.py`. -**Output:** The same file, updated in place with sharding hints, plus YAML config and validation. - -**No separate `_ir.py` file.** Sharding IR is the default path — hints are added directly to the canonical `modeling_*.py`. The legacy pattern of maintaining parallel `modeling_*_ir.py` files is deprecated. - -**Prerequisites:** Familiarity with AD canonical ops (see `ad-model-onboard` skill, Phase 3) and op registration patterns (Phase 4). Refer to the custom op docstrings in `tensorrt_llm/_torch/auto_deploy/custom_ops/` for the complete argument reference (including sharding hints, `tp_mode`, `layer_type`, and which ops accept hints). - -The exported FX graph must fully specify how the model should be sharded: the `apply_sharding_hints` transform combines hints with a runtime `DistConfig` for deterministic, node-local sharding. - -## Step 0 — Sharding-hint delta contract (READ FIRST) - -Adding sharding hints is a **mechanical, structural transform** of the existing `modeling_.py`, NOT a rewrite. The file at the target branch HEAD before your changes is the AUTHORITATIVE source of model logic. - -You MAY introduce ONLY the following changes: - -**ALLOWED:** - -- **A1. Op substitutions:** - - `nn.Linear(...)` / `F.linear(...)` → `torch.ops.auto_deploy.torch_linear_simple(...)` - - `tensor.view(...)` / `tensor.reshape(...)` → `torch.ops.auto_deploy.view(...)` (only when the shape contains a TP-scaled dim) - - `torch.split(...)` / `torch.split_with_sizes(...)` → `torch.ops.auto_deploy.split_with_sizes(...)` -- **A2. Sharding-hint kwargs added** to call sites of: `torch_moe`, `torch_ssm`, `torch_gated_delta_rule`, `torch_causal_conv1d`, `torch_rmsnorm_gated`, `torch_mla`, `torch_attention`, `torch_linear_simple`, `auto_deploy.split_with_sizes`, `auto_deploy.view`. Allowed kwargs: `tp_mode`, `layer_type`, `output_sizes`, `tp_min_local_shape`, `tp_scaled_dim`, `shardable`, `enable_sharding`. -- **A3. Inserting `torch.ops.auto_deploy.all_reduce(..., layer_type=...)`** after rowwise projections / at MoE merge points (single all_reduce after routed + shared sums). -- **A4. Docstring updates:** - - Module-level: a single-line header noting the file uses sharding IR, followed by the existing source-of-truth / HF link block. Example: `"""Llama 3 model (sharding IR)."""`. - - Per-class (MLP, Attention, MoE block, etc.): a short `Sharding strategy:` block listing what each projection maps to (`colwise` / `rowwise` / `all_reduce` / `tp_scaled_dim`). - -**FORBIDDEN (everything else, including but not limited to):** - -- **F1. Replacing ANY `torch.ops.trtllm.*` op with vanilla PyTorch** (e.g. `noaux_tc_op`, `dsv3_router_gemm_op`, fused norm/MLP kernels). The router gate is TP-replicated; there is nothing to shard. AD has no fusion pass that recovers these kernels from a vanilla rewrite — keep the call site verbatim. -- **F2. Changing the input contract** of `forward()` — adding/removing/changing `assert` or `if` statements that change what the caller must pass. -- **F3. Adding/removing/renaming `nn.Module` subclasses, parameters, buffers**, or `register_load_state_dict_pre_hook` registrations. Module hierarchy and state_dict keys must remain identical. -- **F4. Changing dtype handling, scaling factors, normalization order, mask fill values** (e.g. `0.0` vs `-inf` in `masked_fill`), or any other numerical-semantics detail. -- **F5. Renaming methods, changing return types, changing forward signatures**, or reordering operations. -- **F6. "Cleanup" of allegedly unused code paths.** If it is in the file, it stays. -- **F7. Adding code that does not appear in the original** "because a legacy `_ir.py` reference had it" — legacy IR files may be stale or wrong. - -If a change is required that falls outside the allowlist, **STOP and report it to the parent** for explicit human approval BEFORE writing it. Never silently rewrite logic. - -## Reference examples (study before porting) - -The models below already have sharding hints integrated directly into their `modeling_*.py` files. Study them to see how `tp_mode`, `layer_type`, `output_sizes`, `tp_scaled_dim`, `shardable`, `all_reduce`, etc. are placed for different layer types. - -| Model file | Layer types | -|----------|-------------| -| `modeling_nemotron_h.py` | Mamba SSM, MHA, SwiGLU MLP, MoE | -| `modeling_qwen3_5_moe.py` | GatedDeltaNet, Gated MHA, SwiGLU MLP, MoE | -| `modeling_deepseek.py` | MLA, SwiGLU MLP, MoE | -| `modeling_qwen3.py` | MHA, SwiGLU MLP (simplest MHA example) | - -## Step-by-step procedure - -### Step 1: Create a git checkpoint - -Before editing, ensure the file is committed so you can diff against the original: - -```bash -git stash # or commit — ensure a clean baseline to diff against -``` - -### Step 2: Replace linear projections - -For every `self.proj(x)` or `nn.Linear` call, use `torch.ops.auto_deploy.torch_linear_simple` with explicit `tp_mode` and `layer_type`. Always set `tp_mode` unconditionally (no `if _s else "none"`). **Rules:** opening projections (Q/K/V/gate/up/in_proj) → `"colwise"`; closing (O/down/out_proj) → `"rowwise"`; tiny outputs (e.g. `shared_expert_gate` dim 1) → `"none"`; MLA latent projections (q_a, kv_a) → `"none"`. For fused weights split later, pass `output_sizes=[...]`. For GQA, use `tp_min_local_shape=self.head_dim` on K/V colwise lines. - -### Step 3: Replace split / chunk after fused colwise projections - -Use `torch.ops.auto_deploy.split_with_sizes` with `shardable` / `layer_type` where sizes scale with TP. - -### Step 4: Replace view / reshape with concrete head counts - -During `torch.export`, `-1` becomes concrete; after TP, wrong values break. Any reshape whose dimension is a head count that scales with TP must use `torch.ops.auto_deploy.view` with `tp_scaled_dim` set appropriately. Safe cases: flat-to-2D, or `[B,S,-1]` when the input is already correctly sharded. - -### Step 5: Insert `all_reduce` - -After every rowwise projection, add `torch.ops.auto_deploy.all_reduce(..., layer_type=...)`. **Parallel branch rule:** when branches merge by addition, use a **single** `all_reduce` after the sum (e.g. MoE routed + shared expert; parallel attention + MLP residual branches). - -### Step 6: Special ops (Conv1d, SSM, GatedDeltaNet, gated RMSNorm) - -Add sharding hints on `torch_causal_conv1d`, `torch_ssm`, `torch_gated_delta_rule`, `torch_rmsnorm_gated` per docstrings—typically `shardable` / `output_sizes` / `tp_mode` as required. - -### Step 7: MoE - -Pass `layer_type="moe"` into `torch_moe`; `apply_sharding_hints` handles EP/TP. - -### Step 8: Verify registration - -The model's existing registration (`AutoModelForCausalLMFactory.register_custom_model_cls` at the bottom of the file and its import in `__init__.py`) stays unchanged. No new registration is needed — sharding hints do not change the model identity. - -### Step 9: YAML — no per-model opt-in needed - -No YAML change is required to enable the IR path. The default sharding pipeline (`apply_sharding_hints`) auto-detects the presence of `torch.ops.auto_deploy.all_reduce` markers in the exported FX graph and routes IR-marked models to the IR pipeline; non-marked models fall through to the legacy `detect_sharding` + `sharding_transform_executor` pair. The markers you added in Steps 1–7 are sufficient. - -If the model needs a non-default `apply_sharding_hints` config (for example a non-NCCL `allreduce_strategy`, or selective `shard_layers`), add a per-model yaml override under `examples/auto_deploy/model_registry/configs/` that overrides only the keys you need: - -```yaml -transforms: - apply_sharding_hints: - allreduce_strategy: SYMM_MEM - # shard_layers: ['mha', 'mlp'] # optional selective sharding - export_to_gm: - num_moe_experts_for_export: 2 # often required when expert count is large (>64) -``` - -To force the legacy pipeline (e.g. while an IR port has a known bug awaiting fix), add `enable_legacy_sharding.yaml` to the model's `yaml_extra` — that override disables `apply_sharding_hints` and re-enables the legacy stages explicitly. - - -Set `world_size` once, to the **maximum number of GPUs available on the machine**, auto-detected with `python -c 'import torch; print(torch.cuda.device_count())'` (or `nvidia-smi --list-gpus | wc -l`). Do **not** hardcode `world_size: 8` (or any other literal) — porting agents run on heterogeneous hardware and an 8-GPU literal will simply fail to launch on a 2- or 4-GPU machine. If the model's `num_attention_heads` (and, for GQA, `num_key_value_heads`) does not divide the detected GPU count, fall back to the largest power-of-two divisor that does (e.g. 4 on an 8-GPU machine if `num_attention_heads = 12`). Run the end-to-end command exactly once at that size — there is no value in repeating it at multiple smaller sizes, because the offline sharding equivalence test (Step 10b) already exercises 2- and 4-GPU dist configs cheaply. - -Optional `shard_layers` limits which `layer_type` hints are processed; unset means shard all shardable nodes. - -### Step 10a — End-to-end run - -Do not report success until a run completes successfully. - -1. Prefer `python examples/auto_deploy/build_and_run_ad.py --model --use-registry` after updating the registry entry. -2. `apply_sharding_hints` logs should show **`N nodes processed` with N > 0**. -3. If validation fails with infrastructure limits (e.g. head count not divisible by `world_size`), document the assert and compatible sizes; do not "fix" core `sharding.py` / custom op schemas without owner review. -4. If blocked by missing infrastructure support, revert the sharding-hint changes and file a short error report for humans (do not silently patch core transforms). - -**Layer type strings** (for `layer_type` / `shard_layers`): use `"mha"`, `"mla"`, `"mlp"`, `"moe"`, `"ssm"`, `"delta"`, or `"unknown"` (default; skipped when `shard_layers` is set). Match the conventions used in `apply_sharding_hints` and project enums. - -### Step 10b — Sharding equivalence test (MANDATORY) - -Run the offline sharding-IR equivalence test ([`tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_num_correctness.py`](tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_num_correctness.py)) against the modeling file you just edited, under **every** parallelism configuration the test exposes. The port is **not** complete until every configuration passes. Skipping this step or treating a partial pass (e.g. only `tep`) as success is not allowed. - -The test compares a sharded prefill against the unsharded eager reference on a tiny (4-layer, hidden_size=64) instance of the model and asserts `rel_rmse < tol`, where `tol` is the test-defined relative-RMSE tolerance (`REL_RMSE_TOL` constant in [`test_sharding_num_correctness.py`](tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_num_correctness.py); overridable per invocation via the `SHARDING_IR_REL_RMSE_TOL` env var). It uses no PyExecutor / no compile / no checkpoint download, so each cell runs in ~30s on 4xGPU. - -**Run the matrix:** - -```bash -MODEL=tensorrt_llm/_torch/auto_deploy/models/custom/modeling_.py -TEST=tests/unittest/auto_deploy/multigpu/transformations/library/test_sharding_num_correctness.py - -for CFG in tp-only ep-only tep attn-dp; do - pytest "$TEST" --sharding-ir-modeling-file "$MODEL" --sharding-ir-dist-config "$CFG" -s -v \ - 2>&1 | tee /tmp/sharding_ir_${CFG}.log -done -``` - -**Parse the output for each cell. A cell PASSES iff ALL of these are true:** - -1. pytest exit code is `0`. -2. The log contains the line `1 passed` in the pytest summary block. -3. The log contains the rank-0 metrics line `[sharding-ir-eq] |y_s - y_u|: max=... mean=... rel_rmse= (tol=)` and the parsed `rel_rmse` is **strictly less than the parsed `tol`** from the same line. Do not hardcode a tolerance value in the parser — read both `rel_rmse=` and `tol=` from the test's own log and compare them. This stays correct if the test's `REL_RMSE_TOL` is later changed or a per-invocation `SHARDING_IR_REL_RMSE_TOL` is supplied. - -Quick one-liner that prints PASS/FAIL plus the parsed `rel_rmse` and `tol` per cell: - -```bash -for CFG in tp-only ep-only tep attn-dp; do - log=/tmp/sharding_ir_${CFG}.log - if grep -q "1 passed" "$log"; then status=PASS; else status=FAIL; fi - line=$(grep "sharding-ir-eq" "$log" | grep "rel_rmse=" | head -1) - rmse=$(echo "$line" | sed -E 's/.*rel_rmse=([0-9.]+).*/\1/') - tol=$(echo "$line" | sed -E 's/.*\(tol=([0-9.]+)\).*/\1/') - echo "${CFG}: ${status} rel_rmse=${rmse:-NA} tol=${tol:-NA}" -done -``` - -**Failure handling:** - -- A cell failing with `KeyError`, `AttributeError`, `ValueError: You must specify exactly one of input_ids or inputs_embeds`, or any exception *before* `[sharding-ir-eq]` prints means the **modeling code itself** does not yet build / export on a tiny config — fix the modeling code (within the Step 0 allowlist) before proceeding. Do not silently skip the cell. -- A cell where `[sharding-ir-eq]` prints `rel_rmse >= tol` (from the same log line) means a **sharding-hint bug**: a missing `all_reduce`, a wrong `tp_mode`, a `view` without `tp_scaled_dim`, a `split_with_sizes` whose sizes do not scale, etc. Re-read Step 5 (all_reduce), Step 2 (tp_mode), Step 4 (view), Step 3 (split_with_sizes) and the layer-specific patterns. Iterate on the hints until clean. If the failure is small (rel_rmse just slightly above tol) and you have reason to believe it is real numerical noise from the specific layer mix of this model rather than a sharding-hint bug, raise it with the parent agent rather than silently bumping `SHARDING_IR_REL_RMSE_TOL`. -- A cell that the modeling file legitimately does not support (e.g. `ep-only` on a dense model with no MoE) is acceptable only if the failure is a documented `pytest.skip(...)` from the test infrastructure. A silent `FAIL` is **not** acceptable. - -### Step 11 — Pre-finalization self-audit (MANDATORY) - -Before reporting the file as done, you MUST diff your changes against the git baseline: - -```bash -git diff tensorrt_llm/_torch/auto_deploy/models/custom/modeling_.py -``` - -Then classify every hunk into one of the following categories (defined in Step 0): - -| Category | Allowed? | Description | -|---|---|---| -| **A1** | yes | Op substitution (`linear` / `view` / `split`) | -| **A2** | yes | Sharding-hint kwarg added (`tp_mode`, `layer_type`, `output_sizes`, `tp_min_local_shape`, `tp_scaled_dim`, `shardable`, `enable_sharding`) | -| **A3** | yes | `auto_deploy.all_reduce` insertion | -| **A4** | yes | Docstring updates: one-line module header + per-class `Sharding strategy:` blocks | -| **F1** | NO | `torch.ops.trtllm.*` replaced with vanilla PyTorch | -| **F2** | NO | Input contract change (asserts, fallbacks added/removed) | -| **F3** | NO | Module hierarchy / parameter / buffer / load-hook change | -| **F4** | NO | Numerical-semantics change (dtype, scale, mask fill, order) | -| **F5** | NO | Method rename / signature change / op reorder | -| **F6** | NO | Removal of allegedly unused base code | -| **F7** | NO | Code added because a legacy `_ir.py` reference had it (and the base did not) | - -**If you find any F# hunk, REVERT it before reporting done.** Report the full diff classification table back to the parent agent in your final message, with one row per hunk: - -``` -| Hunk lines | Summary of change | Category | Verdict | -|---|---|---|---| -| 234-240 | F.linear → torch_linear_simple, tp_mode="colwise" | A1 + A2 | OK | -| 264-340 | noaux_tc_op replaced with vanilla PyTorch | F1 | REVERTED | -| ... | ... | ... | ... | -``` - -You are NOT done until every row in the table is a yes-allowed category. - -## Layer-specific sharding patterns - -**MHA (standard or gated):** `layer_type="mha"`: q/k/v colwise (GQA: `tp_min_local_shape`), `view` with `tp_scaled_dim` for head dim, o rowwise + `all_reduce`. Fused Q+gate interleaved per head: colwise without `output_sizes`; contiguous Q|K|V fused blocks need `output_sizes`. - -**SwiGLU MLP:** `layer_type="mlp"`: gate/up colwise, down rowwise + `all_reduce`. - -**Mamba / SSM:** `layer_type="ssm"`: in_proj colwise + `output_sizes`, splits shardable, conv1d shardable + `output_sizes`, views, `torch_ssm` shardable, norm gated colwise if weight scales, out rowwise + `all_reduce`. - -**GatedDeltaNet:** `layer_type="delta"`: in_proj_qkv with `output_sizes`, other in_projs colwise, conv1d/splits/views as above, `torch_gated_delta_rule` shardable, out rowwise + `all_reduce`. - -**MoE + shared expert:** `layer_type="moe"`: router replicated; one `all_reduce` after `routed + shared`, not two. - -**MLA (DeepSeek):** `layer_type="mla"`: keep `torch_mla` intact with `shardable=True`—do **not** decompose into separate linears + `torch_attention` (introduces bad `expand`/`view` with concrete head counts). q_a/kv_a latent: `tp_mode="none"`; q_b colwise; `o_proj` rowwise + `all_reduce`. - -**Per-head free Parameters on `torch_attention` (GPT-OSS-style sinks):** when an attention block has a learnable `nn.Parameter` indexed by Q-head count that flows DIRECTLY into `torch_attention` (not through a Linear) — e.g. GPT-OSS's `self.sinks = nn.Parameter(torch.empty(num_heads))` passed as `sinks=self.sinks` — pass `enable_sharding=True` to the `torch_attention(...)` call. The IR's `WeightedParamShardableNode` is registered for `torch_attention` and will slice every direct `get_attr` arg along dim 0 (= head dim) per rank. Q/K/V/O projection weights are unaffected (they belong to the preceding `torch_linear_simple` nodes and are sharded by `LinearShardableNode`). Models with no such head-wise Parameter (qwen3, llama, smollm3, ...) leave `enable_sharding` at its default `False` and the handler no-ops for them. - -## Common pitfalls - -1. **Missing `auto_deploy::view` for head reshapes** — concrete shapes from export break after sharding. -2. **Sharding tiny projections** — dim-1 gates: `tp_mode="none"`. -3. **Double `all_reduce` in MoE** — one merge-point reduction for routed + shared. -4. **Cross-layer parameter contamination** — in `_apply_hint_*` handlers using `get_source_nodes()`, restrict with `allowed_ops` so residual links do not pull weights from other layers. -5. **Missing `num_moe_experts_for_export`** for very large expert counts — export can hang. -6. **Decomposing ops that absorb weights** (e.g. `torch_mla`) — use `shardable` + handler instead of splitting into plain linears. -7. **Interleaved vs contiguous fused weights** — interleaved per-head groups: colwise only; contiguous Q|K|V blocks: require `output_sizes`. -8. **Omitting `layer_type` when using `shard_layers`** — `"unknown"` nodes are skipped; set hints explicitly on sharding-aware ops. -9. **`layer_type` on non-hint ops** — do **not** pass `layer_type` to ops that are not designed for sharding hints (e.g. `torch_l2norm`, `torch_rope_*`); extra positional args break calls. Note: `torch_attention` DOES accept `layer_type` (and `enable_sharding`) — see the per-head Parameters paragraph in "Layer-specific sharding patterns" above. Confirm in `custom_ops/` docstrings which ops accept hints. -10. **Conditional hint values** — no `if _s else "none"`; use unconditional hints and rely on `shard_layers` / transform config. -11. **Replacing `torch.ops.trtllm.*` ops** — `noaux_tc_op`, `dsv3_router_gemm_op`, fused norm/MLP kernels are TP-replicated and must be kept verbatim (rule F1). AD has no fusion pass to recover them from vanilla PyTorch. - -## Validation checklist (human review) - -- All four configurations of the **sharding equivalence test** (Step 10b) pass with the parsed `rel_rmse` strictly below the parsed `tol` from the same rank-0 log line. Report the per-cell `rel_rmse` and `tol` pair. -- `world_size=1`: unsharded path; hints should not break correctness. -- `world_size=`: end-to-end run (Step 10a) at the maximum GPU count auto-detected on the machine (head-divisibility permitting; see Step 10). -- `apply_sharding_hints` node count vs expectation. -- Optional: `shard_layers: ['moe']` to verify selective sharding. diff --git a/.claude/skills/exec-env-check/SKILL.md b/.claude/skills/exec-env-check/SKILL.md new file mode 100644 index 000000000000..3b91216a74b1 --- /dev/null +++ b/.claude/skills/exec-env-check/SKILL.md @@ -0,0 +1,142 @@ +--- +name: exec-env-check +description: >- + Check the local execution environment for GPU availability, Docker support, + and Slurm access. Returns the execution scenario (`satisfied, local, docker`, + `satisfied, local, direct`, `satisfied, slurm, local`, or `not_satisfied`), + the number of available GPUs, and the GPU type. On Slurm login nodes + without local GPUs, the cluster is identified by delegating the hostname + to internal-env-info (hostname-based mode), which owns the + hostname → cluster_name patterns; GPU type and gpus_per_node then come + from that skill's reference files. If internal-env-info is not + installed, the scenario falls back to `not_satisfied` without probing + compute nodes via srun. +tags: [infrastructure, slurm, environment] +license: Apache-2.0 +metadata: + author: NVIDIA Corporation +--- + +# TensorRT-LLM Environment Check + +Detect whether the current machine can run a GPU workload locally (Docker) or via Slurm, and report hardware details. + +## Input + +| Field | Description | Required | +|-------|-------------|----------| +| `required_devices` | Minimum number of GPUs needed | Yes | +| `account` | Slurm account (unused for GPU probing, kept for compatibility) | No | + +## Procedure + +### 1. Check local GPUs + +```bash +timeout 5 nvidia-smi --query-gpu=name,memory.total --format=csv,noheader 2>/dev/null +``` + +If this succeeds: +- Count the number of GPU lines → `available_gpus` +- Extract the GPU name from the first line → `device_type` (e.g., `NVIDIA B200`, `NVIDIA H100 80GB HBM3`) +- Normalize `device_type`: strip `NVIDIA ` prefix and trailing memory info to get the short name (e.g., `B200`, `H100`, `A100`, `L40S`, `RTX 6000`) + +If `nvidia-smi` fails or returns no GPUs → `available_gpus = 0`, `device_type = null`. + +### 2. Check if local GPUs are sufficient + +If `available_gpus >= required_devices`, continue to step 2a to determine whether Docker is available on this host. + +### 2a. Check Docker availability + +```bash +command -v docker >/dev/null 2>&1 && timeout 3 docker info >/dev/null 2>&1 +``` + +- If the command succeeds (Docker CLI exists **and** the daemon responds) → **Result**: `satisfied, local, docker` +- Otherwise (Docker CLI missing, daemon not running, or permission denied) → **Result**: `satisfied, local, direct` + +Include `available_gpus` in the result. Do NOT include `device_type` — local execution does not need it. + +### 3. Check Slurm availability + +If local GPUs are insufficient (or as additional detection), check for Slurm: + +```bash +which squeue 2>/dev/null && squeue --version 2>/dev/null +``` + +If Slurm is NOT available → go to step 5. + +### 4. Resolve GPU type + +**Optional dependency.** Before doing anything else in this step, check whether `skills/internal-env-info/` exists in the toolkit. If it does **not**, skip this step entirely and return `scenario: not_satisfied` with `available_gpus = 0`, `device_type = null`, `gpus_per_node = null`, `cluster_name = null`, `default_models_repo = null`, and `default_user_root_dir = null`. Do **not** report this as an error — the skill is an internal-only dependency. + +When Slurm is available but local `nvidia-smi` returned no GPUs or `device_type` is null (login nodes typically have no GPUs), capture the hostname and delegate cluster identification to `internal-env-info`: + +```bash +hostname -f 2>/dev/null || hostname +``` + +- Pass the captured hostname to `internal-env-info` in **hostname-based mode**. That skill owns the NVIDIA-internal login-host patterns and the hostname → `` mapping; do **not** parse the hostname or hard-code any cluster identifier here. +- It returns the standard output template (`cluster_name`, `device_type`, `gpus_per_node`, `default_models_repo`, `default_user_root_dir`) plus supplementary field (`mfa_style`). +- Set `available_gpus` = `gpus_per_node` (if resolved). +- Set `cluster_name` = the value returned (a placeholder `` token in this skill). The orchestrator uses this to fetch per-cluster info (`mfa_style`, `default_models_repo`, `default_user_root_dir`, `gpus_per_node`) from `internal-env-info`; connection fields (`mounts`, `ssh_host`, `partition`, etc.) come from caller-supplied inputs in `job_spec.json` (with `internal-env-info` default values / ask-the-user fallbacks). +- Set `default_user_root_dir` = the user root directory returned (with `` substituted with the actual SLURM username). Set to `null` if not found. +- If `internal-env-info` returns `null` for `cluster_name` (no pattern matched), set `device_type = null`, `cluster_name = null`, `default_user_root_dir = null` and fall through to Step 5 (`not_satisfied`). + +**Result**: `satisfied, slurm, local` with `device_type`, `gpus_per_node`, `cluster_name`, and `default_user_root_dir`. + +### 5. Not satisfied + +If neither local GPUs nor Slurm is available: +- **Result**: `not_satisfied` + +## Output + +Return a single structured result: + +``` +scenario: +available_gpus: +device_type: +gpus_per_node: +cluster_name: +default_models_repo: +default_user_root_dir: +``` + +**Examples:** + +``` +scenario: satisfied, local, docker +available_gpus: 4 +``` + +``` +scenario: satisfied, local, direct +available_gpus: 4 +``` + +``` +scenario: satisfied, slurm, local +available_gpus: 4 +device_type: B200 +gpus_per_node: 4 +cluster_name: +default_user_root_dir: / +``` + +``` +scenario: not_satisfied +available_gpus: 0 +device_type: null +cluster_name: null +``` + +## Rules + +- Never install drivers or modify the system +- If `nvidia-smi` hangs, use a 5-second timeout: `timeout 5 nvidia-smi ...` +- GPU type is derived from the hostname by matching the cluster name via the `internal-env-info` skill when it is installed — no `srun` allocation needed. If that skill is absent, the scenario falls back to `not_satisfied` (see Step 4); do not error. +- Report the GPU type exactly as found in the mapping (do not guess or fabricate) diff --git a/.claude/skills/exec-local-docker/SKILL.md b/.claude/skills/exec-local-docker/SKILL.md new file mode 100644 index 000000000000..cd73eed6e043 --- /dev/null +++ b/.claude/skills/exec-local-docker/SKILL.md @@ -0,0 +1,121 @@ +--- +name: exec-local-docker +description: >- + Execute a TensorRT-LLM workload locally in Docker. Runs a fully-resolved + Docker command in background, monitors completion, reads logs, and reports + results. Workflow-agnostic — does not need to know if the workload is pytest, + eval, benchmark, or a custom script. +tags: [docker, execution, infrastructure] +license: Apache-2.0 +metadata: + author: NVIDIA Corporation +--- + +# Local Docker Executor + +Run a Docker command locally, monitor it, and report results. + +## Input (from orchestrator prompt) + +The orchestrator passes these fields in the skill prompt: + +| Field | Description | +|-------|-------------| +| `docker_cmd` | Complete `docker run` command string, ready to execute | +| `work_dir` | Local work directory for logs and artifacts | +| `log_file` | Full path to the log file (output redirected here) | +| `model_name` | Short model name for reporting | +| `workflow_type` | `pytest`, `eval`, `custom`, or `benchmark` — for output parsing hints | +| `success_patterns` | Comma-separated patterns indicating success (e.g., `passed,accuracy:`) | +| `failure_patterns` | Comma-separated patterns indicating failure (e.g., `FAILED,Error,AssertionError`) | + +## Procedure + +### Step 0: Resolve Image and Build (when `build_project=true`) + +This executor owns image selection and the build for the local Docker target. Skip the entire step when `build_project=false`. + +1. **Detect the target GPU type.** Use `gpu_type` from `job_spec.json` if upstream env-check resolved it; otherwise probe locally with `nvidia-smi --query-gpu=name --format=csv,noheader | head -1`. +2. **Detect the host CPU arch** with `uname -m` (`x86_64` or `aarch64`). +3. **Resolve the container image.** Read `/jenkins/current_image_tags.properties` and pick the tag whose CPU-arch flavor matches the host. If the orchestrator already passed a `container_image` field in the job spec, use that and skip the lookup. +4. **Map GPU → build arch (`-a` flag):** `H100`/`H200` → `90-real`; `B200`/`GB200`/`B300`/`GB300` → `100-real`; `A100` → `80-real`; `L40S` → `89-real`. Default `100-real` when the GPU is unknown. +5. **Compile.** Invoke the `exec-local-compile` skill with `repo_dir=`, `image=`, `arch=`. Wait for completion. +6. **On failure**, do not launch the workload. Report `BUILD_FAILED` with the last 100 lines of the compile log. + +`build_project`, `gpu_type`, `repo_root`, and (optionally) `container_image` come from `job_spec.json`. + +### Step 1: Launch + +Run the Docker command in background using `run_in_background`: + +```bash + 2>&1 | tee +``` + +Report to the orchestrator: "Launched locally, log at ``" + +### Step 2: Monitor for Hangs + +While waiting for the background process to complete, actively monitor the log +file for hang indicators. Launch a monitoring loop using `run_in_background`: + +```bash +while true; do + sleep 60 + if [ -f "" ] && grep -qi "hang detected" ""; then + echo "HANG_DETECTED: Found 'hang detected' in log file" + docker ps --filter "ancestor=" -q | xargs -r docker kill 2>/dev/null + exit 1 + fi +done +``` + +- If the monitor detects a hang, it kills the Docker container and exits with + code 1. The main background process will also terminate. +- When the main process completes normally (background notification received), + kill the monitoring loop (it is no longer needed). +- If a hang is detected, skip to Step 4 and report `HANG_DETECTED` status + instead of proceeding to normal result collection. + +### Step 3: Wait for Completion + +The Bash tool's `run_in_background` will notify when the process finishes +(either normally or because the container was killed by the hang monitor). + +### Step 4: Read Results + +On completion: + +1. **Read exit code** from the background command result. +2. **Read the last 100 lines** of `` using the Read tool. +3. **If exit code != 0**, also read the first 50 lines to catch early errors (import failures, setup crashes). +4. **Search for patterns**: + - Grep `` for each `success_patterns` entry + - Grep `` for each `failure_patterns` entry + +### Step 5: Report + +Return a structured result: + +``` +Status: PASSED | FAILED | ERROR | HANG_DETECTED +Exit code: +Log file: +Work directory: +Summary: +Errors: +``` + +### Output Parsing by Workflow Type + +- **pytest**: Look for `X passed, Y failed in Zs` summary line +- **eval**: Look for `accuracy:` or `score:` lines; check for `Expected accuracy >= X, but got Y` assertion +- **custom**: No specific patterns — report last 10 lines of output +- **benchmark**: Look for throughput/latency numbers + +## Rules + +- Never run the Docker command in foreground — always use `run_in_background` +- Never `cat` the full log file — use Read with offset/limit or tail +- If the Docker command fails immediately (exit code within seconds), check if the image exists locally +- Report results even if the log file is empty (container may have failed to start) diff --git a/.claude/skills/exec-local-slurm/SKILL.md b/.claude/skills/exec-local-slurm/SKILL.md new file mode 100644 index 000000000000..ecfd3a860fbc --- /dev/null +++ b/.claude/skills/exec-local-slurm/SKILL.md @@ -0,0 +1,443 @@ +--- +name: exec-local-slurm +description: >- + Submit and monitor a Slurm job on a local cluster. Supports two modes: + (1) Persistent allocation (default) — allocates nodes once via nohup salloc, + imports the container once, installs once, and reuses across runs by setting + SLURM env vars and running the sbatch script via bash. (2) One-shot sbatch — + submits a fully-generated Slurm script via sbatch, polls job status, reads + logs on completion, and reports results. Workflow-agnostic — handles pytest, + eval, benchmark, and custom scripts identically. +tags: [slurm, execution, infrastructure] +license: Apache-2.0 +metadata: + author: NVIDIA Corporation +--- + +# Local Slurm Executor + +Submit a Slurm job locally, monitor it, and report results. Uses persistent +allocation by default to eliminate queue wait, container import, and install +overhead on repeated runs. + +## Input (from orchestrator prompt) + +The orchestrator passes these fields in the skill prompt: + +| Field | Description | +|-------|-------------| +| `script_path` | Full local path to the generated `.slurm` script | +| `work_dir` | Local work directory for logs and artifacts | +| `model_name` | Short model name for reporting | +| `workflow_type` | `pytest`, `eval`, `custom`, or `benchmark` — for output parsing hints | +| `success_patterns` | Comma-separated patterns indicating success | +| `failure_patterns` | Comma-separated patterns indicating failure | +| `log_file_pattern` | Log filename pattern with `%j` placeholder (e.g., `llama_auto_test_%j.out`) | +| `persistent_mode` | Default: `true` for all local slurm workflows. Set to `false` to force one-shot sbatch (opt-out). | +| `release_allocation` | `true` to release the current allocation and stop. Only set when the user explicitly says no more jobs are needed. Default: `false`. Never auto-release. | +| `alloc_time_limit` | Walltime for the persistent allocation. Default: `04:00:00`. | +| `docker_image` | Container image for persistent container import. From `job_spec.json`. | +| `container_name` | Container name used in the `.slurm` script (e.g., `llama_auto_test`). Must match exactly. From `job_spec.json`. | +| `container_mounts` | Comma-separated mount mappings. From `job_spec.json`. | +| `repo_root` | Repo root path (for locating state file and project path). | +| `slurm_params` | Slurm parameters object: `partition`, `account`, `nodes`, `ntasks`, `ntasks_per_node`, `gpus_per_node`. From `job_spec.json`. | + +## Procedure + +### Pre-step: Resolve Image and Build (when `build_project=true`) + +This executor owns image selection and the build for the local SLURM cluster. Skip the entire pre-step when `build_project=false`. + +1. **Read the cluster GPU type** from `device_type` in `job_spec.json` (resolved upstream by env-check / internal-env-info). If absent, probe from any compute node via `srun -p --ntasks-per-node=1 nvidia-smi --query-gpu=name --format=csv,noheader | head -1`. If `skills/internal-env-info/` is not installed, skip the upstream lookup silently and rely on the srun probe — do not report the missing skill as an error. +2. **Determine the partition's CPU arch.** Look up `job_spec.slurm_env.partitions[]` for the entry where `name == partition` and read its `arch` (`x86_64` / `aarch64`). Case-executor's Step 2.5 already detected this — do **not** run `scontrol show node` here. If `slurm_env` is absent or the entry's `arch` is null (case-executor ran without SLURM tools), fall back to inferring from `device_type`: Grace-based parts (`GB*` / `GH*`) → `aarch64`; everything else → `x86_64`. +3. **Resolve the container image.** If `docker_image` is already set in `job_spec.json`, use it directly. Otherwise read `/jenkins/current_image_tags.properties` and pick the tag matching the partition's CPU arch. +4. **Map GPU → build arch (`-a` flag):** `H100`/`H200` → `90-real`; `B200`/`GB200`/`B300`/`GB300` → `100-real`; `A100` → `80-real`; `L40S` → `89-real`. Default `100-real`. +5. **Compile.** Invoke the `exec-slurm-compile` skill with `repo_dir=`, `partition`, `account`, `container_image=`, `user_root_dir`, `arch`. Wait for completion. +6. **On failure**, do not allocate or run the workload. Report `BUILD_FAILED` with the last 100 lines of the build log. + +`build_project`, `device_type`, `repo_root`, `partition`, `account`, `user_root_dir`, and (optionally) `docker_image` come from `job_spec.json`. + +### Step 0: Allocation Management + +This step runs before any execution. It determines whether to reuse an existing +persistent allocation, create a new one, or fall through to one-shot sbatch. + +#### Step 0A — Release mode + +If `release_allocation=true`: + +1. Read `/work_dirs/.slurm_alloc.json` +2. If file exists and `job_id` is present: + ```bash + scancel + ``` +3. Delete the state file and `.salloc.log` +4. Report: "Allocation released." Stop. + +#### Step 0B — One-shot mode + +If `persistent_mode=false`, skip entirely to Step 1 (One-Shot Path). + +#### Step 0C — Validate existing allocation + +Read `/work_dirs/.slurm_alloc.json`. + +- If state file **does not exist** → go to Step 0D. +- If state file exists, validate the allocation: + ```bash + squeue -j -h -o "%T %L" + ``` + - **RUNNING + remaining > 5 min + params compatible** → **reuse** (skip to Step 0F) + - **RUNNING + remaining <= 5 min** → warn user ("Allocation expiring soon"), `scancel `, delete state file, go to Step 0D + - **RUNNING + params incompatible** → `scancel `, delete state file, go to Step 0D + - **Empty output / error** (job gone) → stale state file, delete it, go to Step 0D + +**Params compatibility check:** +- `partition` must match +- `nodes` in state must be **>=** requested nodes (a 2-node allocation can serve 1-node jobs) +- `docker_image` must match + +#### Step 0D — Justify and prepare allocation + +Before allocating, ensure no orphaned allocations exist and log the reason: + +1. Check for existing persistent jobs: + ```bash + squeue -u $(whoami) -h -o "%i %T %j" --name=-trtllm.persistent + ``` +2. If found → orphan (state file was missing/corrupt). Cancel it: + ```bash + scancel + ``` + Log: "Released orphaned allocation — state file was missing." +3. Log justification: "No reusable allocation found. Allocating node(s) on partition for ." + +#### Step 0E — Allocate new + +Convert `alloc_time_limit` from `HH:MM:SS` to seconds (e.g., `04:00:00` → `14400`). + +```bash +mkdir -p /work_dirs +nohup salloc --partition= --account= \ + --nodes= --time= \ + --job-name=-trtllm.persistent \ + sleep \ + > /work_dirs/.salloc.log 2>&1 & +``` + +Retrieve job ID: +```bash +squeue -u $(whoami) -h -o "%i" --name=-trtllm.persistent +``` + +If no job appears after 10 seconds, read `/work_dirs/.salloc.log` +for errors (bad partition, invalid account, etc.) and report to user. + +Otherwise, poll until RUNNING (every 10s, max 60 polls). Once RUNNING, get +the nodelist: +```bash +squeue -j -h -o "%N" +``` + +**Import container** on all nodes (using the **same container name** as the +`.slurm` script — critical for pyxis reuse): +```bash +srun --jobid= -N --ntasks-per-node=1 \ + --container-image= \ + --container-name= true +``` + +**Warm up filesystem mounts** on all nodes — `ls` each mounted path so +Lustre/NFS metadata is cached for later use: +```bash +srun --jobid= -N --ntasks-per-node=1 \ + --container-name= \ + --container-mounts= \ + bash -c 'for p in ...; do ls "$p" > /dev/null 2>&1; done' +``` +Parse mount targets from `container_mounts` — the right-hand side of each +`host:container` pair. + +**Check GPU status** on all nodes (no container needed — `nvidia-smi` is on +the host): +```bash +srun --jobid= -N --ntasks-per-node=1 \ + bash -c 'echo "=== $(hostname) ===" && nvidia-smi --query-compute-apps=pid,name,used_memory --format=csv,noheader' +``` +- If output shows no processes → GPUs are clean, proceed +- If unexpected processes found → warn user: "GPU processes found on + : . These may interfere with the job." + +**Run install** on all nodes with `--container-writable` (so packages persist +in the named container across srun calls). Install all common requirements +upfront so that subsequent jobs of any workflow type can skip install: +```bash +srun --jobid= -N --ntasks-per-node=1 \ + --container-name= --container-writable \ + --container-mounts= \ + bash -c 'cd && pip install -e . && pip install -r requirements-dev.txt && \ + if [ -f examples/trtllm-eval/requirements.txt ]; then pip install -r examples/trtllm-eval/requirements.txt; fi' +``` +For custom workflow with `skip_install=true`, skip this step entirely. + +**Write state file** `/work_dirs/.slurm_alloc.json`: +```json +{ + "job_id": "", + "container_name": "", + "nodelist": "", + "partition": "", + "account": "", + "nodes": , + "gpus_per_node": , + "docker_image": "", + "container_mounts": "", + "allocated_at": "", + "time_limit": "", + "installed": true +} +``` + +Proceed to Step 1 (Persistent Path). + +#### Step 0F — Reuse path validation + +Allocation is valid. Check if container name matches current request: + +- If `container_name` in state file **matches** `container_name` from + `job_spec.json` → proceed to Step 1 (Persistent Path). Container and install + are already set up. +- If **different** (model changed) → import the new container, warm up mounts, + and install: + ```bash + srun --jobid= -N --ntasks-per-node=1 \ + --container-image= \ + --container-name= true + srun --jobid= -N --ntasks-per-node=1 \ + --container-name= \ + --container-mounts= \ + bash -c 'for p in ; do ls "$p" > /dev/null 2>&1; done' + srun --jobid= -N --ntasks-per-node=1 \ + --container-name= --container-writable \ + --container-mounts= \ + bash -c 'cd && pip install -e . && pip install -r requirements-dev.txt' + ``` + Update `container_name` in state file. Proceed to Step 1 (Persistent Path). + +--- + +### Step 1: Execute + +Two execution paths depending on `persistent_mode`. + +#### Persistent Path (persistent_mode=true) + +**Pre-flight time check** — verify the allocation has enough remaining time: +```bash +squeue -j -h -o "%L" +``` +If remaining time < job's `time_limit` from `slurm_params`, warn user: +"Allocation has left but job expects . The job may +be killed early." + +**Pre-flight GPU check** — verify GPUs are not occupied by leftover processes: +```bash +srun --jobid= -N --ntasks-per-node=1 \ + bash -c 'procs=$(nvidia-smi --query-compute-apps=pid,name,used_memory --format=csv,noheader 2>/dev/null); [ -n "$procs" ] && echo "WARNING: GPU processes on $(hostname): $procs"' +``` +If unexpected processes found, warn user before proceeding. + +**Run the sbatch script** with SLURM env vars — the `#SBATCH` directives are +comments when run via `bash`; the inner `srun` inherits the env vars and uses +the persistent allocation: +```bash +export SLURM_JOB_ID= +export SLURM_JOB_NUM_NODES= +export SLURM_NNODES= +export SLURM_NTASKS= +export SLURM_NTASKS_PER_NODE= +export SLURM_NODELIST= +export SLURM_JOB_NODELIST= +bash 2>&1 | tee /_.log +``` + +Run with `run_in_background`. Container reuse is automatic (pyxis skips +re-import when `--container-name` already exists on the node). The install +step (Step 1 srun) always runs to keep the container up to date. + +**Step cancellation:** If a job hangs or the user wants to abort, cancel just +the srun step without killing the allocation: +```bash +# List active steps: +squeue -s -j +# Cancel a specific step (e.g., step 0): +scancel .0 +``` +The allocation stays RUNNING — new jobs can run immediately. The job name +`-trtllm.persistent` uniquely identifies allocations created by this +skill via `squeue --name=-trtllm.persistent`. + +#### One-Shot Path (persistent_mode=false) + +```bash +sbatch +``` + +Parse the job ID from output: `Submitted batch job `. + +If sbatch fails, report the error immediately and stop. + +### Step 2: Report Submission + +**One-shot mode:** +``` +Job ID: +Script: +Work directory: +Log files: / +Monitor: squeue -j +``` + +**Persistent mode:** +``` +Allocation: (persistent, remaining: