Hi,

I am testing "DeepSeek-R1-Distill-Qwen-32B" in SageMaker and I want to increase MAX_TOTAL_TOKENS to a 6k+. I am getting the following error. Any tips on how to fix this?

ValueError: No cached version found for deepseek-ai/DeepSeek-R1-Distill-Qwen-32B with {'task': 'text-generation', 'batch_size': 8, 'num_cores': 8, 'auto_cast_type': 'bf16', 'sequence_length': 8192, 'compiler_type': 'neuronx-cc', 'compiler_version': '2.15.143.0+e39249ad', 'checkpoint_id': 'deepseek-ai/DeepSeek-R1-Distill-Qwen-32B', 'checkpoint_revision': 'd66bcfc2f3fd52799f95943264f32ba15ca0003d'}.You can start a discussion to request it on https://huggingface.co/aws-neuron/optimum-neuron-cacheAlternatively, you can export your own neuron model as explained in https://huggingface.co/docs/optimum-neuron/main/en/guides/export_model#exporting-neuron-models-using-neuronx-tgi

Here's my code:

import json
import sagemaker
import boto3
from sagemaker.huggingface import HuggingFaceModel, get_huggingface_llm_image_uri

try:
    role = sagemaker.get_execution_role()
except ValueError:
    iam = boto3.client("iam")
    role = iam.get_role(RoleName="sagemaker_execution_role")["Role"]["Arn"]

# Hub Model configuration. https://huggingface.co/models
hub = {
    "HF_MODEL_ID": "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B",
    "HF_NUM_CORES": "8",
    "HF_AUTO_CAST_TYPE": "bf16",
    "MAX_BATCH_SIZE": "8",
    "MAX_INPUT_TOKENS": "7000",
    "MAX_TOTAL_TOKENS": "8192",
}


region = boto3.Session().region_name
image_uri = f"763104351884.dkr.ecr.{region}.amazonaws.com/huggingface-pytorch-tgi-inference:2.1.2-optimum0.0.27-neuronx-py310-ubuntu22.04"

# create Hugging Face Model Class
huggingface_model = HuggingFaceModel(
    image_uri=image_uri,
    env=hub,
    role=role,
)

# deploy model to SageMaker Inference
predictor = huggingface_model.deploy(
    initial_instance_count=1,
    instance_type="ml.inf2.24xlarge",
    container_startup_health_check_timeout=1800,
    volume_size=512,
)

# send request
predictor.predict(
    {
        "inputs": "What is is the capital of France?",
        "parameters": {
            "do_sample": True,
            "max_new_tokens": 128,
            "temperature": 0.7,
            "top_k": 50,
            "top_p": 0.95,
        }
    }
)