Add dependencies check for ptq

Signed-off-by: Kai Xu <kaix@nvidia.com>

Author: kaix-nv

.claude/skills/ptq/SKILL.md

@@ -24,6 +24,24 @@ Check the support table in `examples/llm_ptq/README.md` for verified HF models.
 - **Listed** → supported, use `hf_ptq.py` (step 4A/4B)
 - **Not listed** → read `references/unsupported-models.md` to determine if `hf_ptq.py` can still work or if a custom script is needed (step 4C)
 
+## Step 2.5 — Check for model-specific dependencies
+
+If the model uses `trust_remote_code` (check `config.json` for `auto_map`), inspect its custom Python files for imports not present in the container:
+
+```bash
+grep -h "^from \|^import " <model_path>/modeling_*.py | sort -u
+```
+
+**Known dependency patterns:**
+
+| Import found | Packages to install |
+| --- | --- |
+| `from mamba_ssm` / `from causal_conv1d` | `mamba-ssm causal-conv1d` (Mamba/hybrid models: NemotronH, Jamba) |
+
+If extra deps are needed:
+- **Launcher (4B)**: set `EXTRA_PIP_DEPS` in the task's `environment` section — `ptq.sh` installs them automatically
+- **Manual (4A)**: `unset PIP_CONSTRAINT && pip install <deps>` before running `hf_ptq.py`
+
 ## Step 3 — Choose quantization format
 
 **First**, check for a model-specific recipe:
@@ -124,6 +142,7 @@ Report the path and size to the user.
 
 ## Common Pitfalls
 
+- **Model-specific dependencies**: Models with `trust_remote_code` may import packages not in the container (e.g., `mamba-ssm` for hybrid Mamba models). See Step 2.5. Use `EXTRA_PIP_DEPS` env var with the launcher, or install manually before running `hf_ptq.py`
 - **Transformers version**: New models may need a newer version of transformers than what's installed. Check `config.json` for `transformers_version`. In containers, beware of `PIP_CONSTRAINT` blocking upgrades — see `references/slurm-setup-ptq.md` for workarounds
 - **Gated datasets**: Some calibration datasets require HF authentication. Ensure `HF_TOKEN` is set in the job environment, or use `--dataset cnn_dailymail` as a non-gated alternative
 - **NFS root_squash + Docker**: See `skills/common/slurm-setup.md` section 5

tools/launcher/common/hf/ptq.sh

@@ -25,6 +25,13 @@
 
 set -e
 
+# Install extra pip dependencies if specified (e.g., mamba-ssm for hybrid Mamba models).
+if [ -n "$EXTRA_PIP_DEPS" ]; then
+    echo "Installing extra dependencies: $EXTRA_PIP_DEPS"
+    unset PIP_CONSTRAINT
+    pip install $EXTRA_PIP_DEPS
+fi
+
 REPO=""
 LOCAL_DIR=""
 PTQ_ARGS=()