fix: export review feedback - dtype mismatch, default, validation

- Update torch_dtype in config.json when export(dtype=...) overrides it
- Remove hardcoded num_target_layers=36 default (use base_config directly)
- Add state dict validation (assert non-empty after extraction)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: Chenhan Yu <chenhany@nvidia.com>

Author: ChenhanYu

modelopt/torch/export/plugins/hf_spec_export.py

@@ -340,7 +340,7 @@ def _export_config(self):
             "torch_dtype": str(getattr(base_config, "torch_dtype", torch.bfloat16)).replace(
                 "torch.", ""
             ),
-            "num_target_layers": getattr(base_config, "num_hidden_layers", 36),
+            "num_target_layers": base_config.num_hidden_layers,
         }
 
         # Add layer_types if present (Qwen3-style)
@@ -366,12 +366,16 @@ def export(self, export_dir: Path | str, dtype: torch.dtype | None = None):
 
         # Export state dict
         drafter_sd = self._extract_state_dict(full_sd)
+        assert drafter_sd, "No dflash_module weights found in state dict"
         if dtype is not None and hf_quant_config is None:
             drafter_sd = {k: v.to(dtype) for k, v in drafter_sd.items()}
         save_file(drafter_sd, f"{export_dir}/model.safetensors")
 
         # Export config
         drafter_config = self._export_config()
+        # Update torch_dtype to match actual exported weights
+        if dtype is not None:
+            drafter_config["torch_dtype"] = str(dtype).replace("torch.", "")
         if hf_quant_config is not None:
             drafter_config["quantization_config"] = hf_quant_config
         with open(f"{export_dir}/config.json", "w") as f: