dreamzero / scripts /inference /build_trt_engine.sh

Upload dreamzero source code (no model checkpoints)

fbd9366 verified 3 months ago

4.88 kB

#!/usr/bin/env bash
# Build a TensorRT engine from a DreamZero checkpoint.
#
# Usage (recommended — with real calibration data):
#   bash scripts/inference/build_trt_engine.sh \
#       --model-path ./checkpoints/DreamZero-DROID \
#       --tensorrt nvfp4 \
#       --dataset-path ./data/droid_lerobot \
#       --cuda-device 0
#
# Usage (without dataset — acceptable for fp16, not recommended for nvfp4/fp8):
#   bash scripts/inference/build_trt_engine.sh \
#       --model-path ./checkpoints/DreamZero-DROID \
#       --tensorrt nvfp4 \
#       --cuda-device 0
#
# The engine is saved to:
#   {model_path}/tensorrt/wan/WanModel_{precision}.trt
#
# Supported precisions: nvfp4 (recommended), fp8, fp16
#
# For quantized precisions (nvfp4, fp8), ModelOpt calibrates quantization
# parameters using real forward passes.  Providing --dataset-path is strongly
# recommended — random dummy inputs are used as fallback but reduce accuracy.
#
# ENABLE_TENSORRT=true must be set before any groot modules are imported
# (it controls flash-attention compatibility mode for ONNX/TRT export).
# This script sets it and launches the Python build script via torchrun so
# that RANK / WORLD_SIZE env vars are available for GrootSimPolicy init.

# export HF_HUB_CACHE=/mnt/aws-lfs-02/shared/ckpts
set -euo pipefail

MODEL_PATH=""
TENSORRT_PRECISION=""
CUDA_DEVICE="0"
DATASET_PATH=""
NUM_CALIBRATION_TRAJS="2"

while [[ $# -gt 0 ]]; do
    case $1 in
        --model-path)
            MODEL_PATH="$2"
            shift 2
            ;;
        --tensorrt)
            TENSORRT_PRECISION="$2"
            shift 2
            ;;
        --cuda-device)
            CUDA_DEVICE="$2"
            shift 2
            ;;
        --dataset-path)
            DATASET_PATH="$2"
            shift 2
            ;;
        --num-calibration-trajs)
            NUM_CALIBRATION_TRAJS="$2"
            shift 2
            ;;
        -h|--help)
            echo "Usage: $0 --model-path <path> --tensorrt <precision> [OPTIONS]"
            echo ""
            echo "Options:"
            echo "  --model-path PATH             Path to DreamZero checkpoint directory"
            echo "  --tensorrt PRECISION          TRT precision: nvfp4 (recommended), fp8, fp16"
            echo "  --dataset-path PATH           LeRobot dataset for real calibration (recommended for nvfp4/fp8)"
            echo "  --num-calibration-trajs N     Number of calibration trajectories (default: 2)"
            echo "  --cuda-device ID              CUDA device index (default: 0)"
            exit 0
            ;;
        *)
            echo "Unknown argument: $1" >&2
            exit 1
            ;;
    esac
done

if [[ -z "$MODEL_PATH" ]]; then
    echo "Error: --model-path is required" >&2
    exit 1
fi

if [[ -z "$TENSORRT_PRECISION" ]]; then
    echo "Error: --tensorrt is required (e.g. nvfp4, fp8, fp16)" >&2
    exit 1
fi

if [[ ! -d "$MODEL_PATH" ]]; then
    echo "Error: checkpoint directory not found: $MODEL_PATH" >&2
    exit 1
fi

ENGINE_PATH="${MODEL_PATH}/tensorrt/wan/WanModel_${TENSORRT_PRECISION}.trt"

echo "=========================================="
echo "DreamZero TensorRT Engine Builder"
echo "  Checkpoint         : $MODEL_PATH"
echo "  Precision          : $TENSORRT_PRECISION"
echo "  CUDA device        : $CUDA_DEVICE"
echo "  Dataset (calibrate): ${DATASET_PATH:-<none — using dummy inputs>}"
echo "  Calibration trajs  : $NUM_CALIBRATION_TRAJS"
echo "  Output             : $ENGINE_PATH"
echo "=========================================="

# ENABLE_TENSORRT must be set before Python imports any groot model modules
# (it activates flash-attention compatibility mode required for ONNX/TRT export).
export ENABLE_TENSORRT=true
export CUDA_VISIBLE_DEVICES="$CUDA_DEVICE"
export ATTENTION_BACKEND="TE"
export HYDRA_FULL_ERROR=1

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"

# Build the Python argument list.
PYTHON_ARGS=(
    --model-path "$MODEL_PATH"
    --tensorrt   "$TENSORRT_PRECISION"
    --num-calibration-trajs "$NUM_CALIBRATION_TRAJS"
)
if [[ -n "$DATASET_PATH" ]]; then
    PYTHON_ARGS+=(--dataset-path "$DATASET_PATH")
fi

# torchrun sets RANK / WORLD_SIZE / MASTER_ADDR / MASTER_PORT which are
# required by GrootSimPolicy's distributed init.
torchrun \
    --standalone \
    --nproc_per_node=1 \
    "${REPO_ROOT}/scripts/inference/build_trt_engine_droid.py" \
    "${PYTHON_ARGS[@]}"

echo "=========================================="
echo "Engine built successfully: $ENGINE_PATH"
echo ""
echo "Run inference with:"
echo "  CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.run --standalone --nproc_per_node=2 \\"
echo "      socket_test_optimized_AR.py --port 5000 --enable-dit-cache \\"
echo "      --model-path ${MODEL_PATH} --tensorrt ${TENSORRT_PRECISION}"
echo "=========================================="

	#!/usr/bin/env bash
	# Build a TensorRT engine from a DreamZero checkpoint.
	#
	# Usage (recommended — with real calibration data):
	# bash scripts/inference/build_trt_engine.sh \
	# --model-path ./checkpoints/DreamZero-DROID \
	# --tensorrt nvfp4 \
	# --dataset-path ./data/droid_lerobot \
	# --cuda-device 0
	#
	# Usage (without dataset — acceptable for fp16, not recommended for nvfp4/fp8):
	# bash scripts/inference/build_trt_engine.sh \
	# --model-path ./checkpoints/DreamZero-DROID \
	# --tensorrt nvfp4 \
	# --cuda-device 0
	#
	# The engine is saved to:
	# {model_path}/tensorrt/wan/WanModel_{precision}.trt
	#
	# Supported precisions: nvfp4 (recommended), fp8, fp16
	#
	# For quantized precisions (nvfp4, fp8), ModelOpt calibrates quantization
	# parameters using real forward passes. Providing --dataset-path is strongly
	# recommended — random dummy inputs are used as fallback but reduce accuracy.
	#
	# ENABLE_TENSORRT=true must be set before any groot modules are imported
	# (it controls flash-attention compatibility mode for ONNX/TRT export).
	# This script sets it and launches the Python build script via torchrun so
	# that RANK / WORLD_SIZE env vars are available for GrootSimPolicy init.

	# export HF_HUB_CACHE=/mnt/aws-lfs-02/shared/ckpts
	set -euo pipefail

	MODEL_PATH=""
	TENSORRT_PRECISION=""
	CUDA_DEVICE="0"
	DATASET_PATH=""
	NUM_CALIBRATION_TRAJS="2"

	while [[ $# -gt 0 ]]; do
	case $1 in
	--model-path)
	MODEL_PATH="$2"
	shift 2
	;;
	--tensorrt)
	TENSORRT_PRECISION="$2"
	shift 2
	;;
	--cuda-device)
	CUDA_DEVICE="$2"
	shift 2
	;;
	--dataset-path)
	DATASET_PATH="$2"
	shift 2
	;;
	--num-calibration-trajs)
	NUM_CALIBRATION_TRAJS="$2"
	shift 2
	;;
	-h\|--help)
	echo "Usage: $0 --model-path <path> --tensorrt <precision> [OPTIONS]"
	echo ""
	echo "Options:"
	echo " --model-path PATH Path to DreamZero checkpoint directory"
	echo " --tensorrt PRECISION TRT precision: nvfp4 (recommended), fp8, fp16"
	echo " --dataset-path PATH LeRobot dataset for real calibration (recommended for nvfp4/fp8)"
	echo " --num-calibration-trajs N Number of calibration trajectories (default: 2)"
	echo " --cuda-device ID CUDA device index (default: 0)"
	exit 0
	;;
	*)
	echo "Unknown argument: $1" >&2
	exit 1
	;;
	esac
	done

	if [[ -z "$MODEL_PATH" ]]; then
	echo "Error: --model-path is required" >&2
	exit 1
	fi

	if [[ -z "$TENSORRT_PRECISION" ]]; then
	echo "Error: --tensorrt is required (e.g. nvfp4, fp8, fp16)" >&2
	exit 1
	fi

	if [[ ! -d "$MODEL_PATH" ]]; then
	echo "Error: checkpoint directory not found: $MODEL_PATH" >&2
	exit 1
	fi

	ENGINE_PATH="${MODEL_PATH}/tensorrt/wan/WanModel_${TENSORRT_PRECISION}.trt"

	echo "=========================================="
	echo "DreamZero TensorRT Engine Builder"
	echo " Checkpoint : $MODEL_PATH"
	echo " Precision : $TENSORRT_PRECISION"
	echo " CUDA device : $CUDA_DEVICE"
	echo " Dataset (calibrate): ${DATASET_PATH:-<none — using dummy inputs>}"
	echo " Calibration trajs : $NUM_CALIBRATION_TRAJS"
	echo " Output : $ENGINE_PATH"
	echo "=========================================="

	# ENABLE_TENSORRT must be set before Python imports any groot model modules
	# (it activates flash-attention compatibility mode required for ONNX/TRT export).
	export ENABLE_TENSORRT=true
	export CUDA_VISIBLE_DEVICES="$CUDA_DEVICE"
	export ATTENTION_BACKEND="TE"
	export HYDRA_FULL_ERROR=1

	SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
	REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"

	# Build the Python argument list.
	PYTHON_ARGS=(
	--model-path "$MODEL_PATH"
	--tensorrt "$TENSORRT_PRECISION"
	--num-calibration-trajs "$NUM_CALIBRATION_TRAJS"
	)
	if [[ -n "$DATASET_PATH" ]]; then
	PYTHON_ARGS+=(--dataset-path "$DATASET_PATH")
	fi

	# torchrun sets RANK / WORLD_SIZE / MASTER_ADDR / MASTER_PORT which are
	# required by GrootSimPolicy's distributed init.
	torchrun \
	--standalone \
	--nproc_per_node=1 \
	"${REPO_ROOT}/scripts/inference/build_trt_engine_droid.py" \
	"${PYTHON_ARGS[@]}"

	echo "=========================================="
	echo "Engine built successfully: $ENGINE_PATH"
	echo ""
	echo "Run inference with:"
	echo " CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.run --standalone --nproc_per_node=2 \\"
	echo " socket_test_optimized_AR.py --port 5000 --enable-dit-cache \\"
	echo " --model-path ${MODEL_PATH} --tensorrt ${TENSORRT_PRECISION}"
	echo "=========================================="