---

name: inference-export description: "Export Composer or PyTorch models for inference with TorchScript or ONNX, including checkpoint loading, export callbacks, and HuggingFace integration caveats." metadata: disco-role: operating disable-model-invocation: true license: Apache 2.0

Composer Inference Export

Use this sub-skill when a user asks to make a Composer or PyTorch model ready for inference through TorchScript, ONNX, checkpoint-backed export, ExportForInferenceCallback, or optional HuggingFace/PEFT integration.

Route here

• Export a model to TorchScript or ONNX.
• Export from an existing checkpoint with load_path and optional load_object_store.
• Attach ExportForInferenceCallback so export happens at the end of training.
• Use Trainer.export_for_inference from an already-constructed Trainer.
• Choose sample_input, dynamic_axes, input_names, and output_names.
• Apply surgery_algs, transforms, or quantize_dynamic before export.
• Validate a saved TorchScript or ONNX artifact.
• Handle HuggingFaceModel tokenizer, embedding, PEFT, or checkpoint metadata caveats.

Reroute

• Training, checkpoint saving, or resume flow details: use ../training/SKILL.md.
• Logger destinations, object-store upload, and artifact routing beyond export arguments: use ../observability/SKILL.md.
• Distributed launch, wrapping, FSDP setup, or backend rank topology: use ../distributed/SKILL.md.
• Speedup method choice before export: use ../methods/SKILL.md, then return here for artifact generation.

Read first

• Export and inference: recipes, signatures, arguments, and validation checks.
• HuggingFace integration: optional mosaicml[nlp], PEFT, tokenizer/config/embedding caveats, and checkpoint metadata.
• Troubleshooting: export-specific failure modes and recovery steps.
• TorchScript export smoke: a safe no-download export/load check.

Core export choices

1. Choose save_format="torchscript" for a self-contained scripted/traced artifact.
2. Choose save_format="onnx" when the target runtime expects ONNX and you can provide a representative sample_input.
3. Use load_path when the model object needs weights from a Composer checkpoint before export.
4. Use surgery_algs for pre-load model surgery and transforms for post-load export-time transforms.
5. Use quantize_dynamic only when the target model layers and runtime support dynamic quantization.
6. Validate the output artifact immediately on the same example input.

Common API skeleton

from composer.utils import export_for_inference, quantize_dynamic

export_for_inference(
    model=model,
    save_format="torchscript",
    save_path="model.pt",
    sample_input=example_batch,
    transforms=[quantize_dynamic],
)

For ONNX, provide sample_input and usually explicit names or dynamic axes:

export_for_inference(
    model=model,
    save_format="onnx",
    save_path="model.onnx",
    sample_input=(inputs,),
    input_names=["input"],
    output_names=["logits"],
    dynamic_axes={"input": {0: "batch"}, "logits": {0: "batch"}},
)

Callback pattern

Use ExportForInferenceCallback when export should be part of the training lifecycle:

from composer.callbacks import ExportForInferenceCallback
from composer import Trainer

trainer = Trainer(
    model=model,
    train_dataloader=train_loader,
    max_duration="1ep",
    callbacks=[ExportForInferenceCallback(save_format="torchscript", save_path="model.pt")],
)
trainer.fit()

For ONNX callbacks, pass sample_input or allow the callback to capture a batch after the dataloader runs.

Distributed and checkpoint caveats

• Only global rank zero writes the export artifact.
• If the model is wrapped in DDP, export the underlying module rather than the wrapper.
• Do not export an FSDP-wrapped module directly; recreate/load the model without FSDP wrapping first.
• load_strict=False is the export utility default; choose strictness intentionally when loading changed heads or adapter-only weights.
• Remote checkpoint paths need matching object-store or logger configuration from the observability route.

HuggingFace quick checks

• Install the NLP extra before using HuggingFaceModel.
• Pass the tokenizer when checkpoint metadata should preserve tokenizer files.
• If tokenizer vocabulary is larger than model embeddings, either resize embeddings before wrapping or set allow_embedding_resizing=True.
• PEFT support is optional and Composer's wrapper expects supported PEFT configuration, with LoRA as the intended adapter path.

Bundled smoke script

Run from this sub-skill directory:

python scripts/export_torchscript_smoke.py

The script exports a tiny CPU model to a temporary TorchScript file, reloads it, and compares outputs.

Ask or stop before proceeding

• The export target is ONNX but onnx/onnxruntime are not available and validation is required.
• The model requires downloading remote weights/tokenizers and network approval is unclear.
• The checkpoint is untrusted, private, remote, or requires credentials.
• The model is FSDP-wrapped and the unwrapped construction path is unknown.
• The target runtime requires shapes, names, quantization, or operator sets that the source model may not support.