#2783·text-generation-inference

"sharded is not supported for AutoModel" Error When Deploying SageMaker Endpoint For Qwen 2.5 7B Trained via SageMaker

Author: jjbuckCreated Nov 26, 2024Updated Sep 23, 2025

System Info

Docker image: 763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-tgi-inference:2.1.1-tgi1.4.0-gpu-py310-cu121-ubuntu20.04
Deployment: SageMaker endpoint, using the HuggingFaceModel object.

Information

Docker
The CLI directly

Tasks

An officially supported command
My own modifications

Reproduction

I'm trying to deploy a SageMaker endpoint after having fine-tuned a model (Qwen 2.5 7B) using the official SageMaker HuggingFace estimator. The error I'm running into is NotImplementedError: sharded is not supported for AutoModel. The full stack trace is reproduced at the end of this post.

For context, here's how I create the SageMaker training job.

huggingface_estimator = HuggingFace(
    entry_point          = 'train.py',       
    dependencies=['requirements.txt'],         
    source_dir           = './',             
    instance_type        = 'ml.p4d.24xlarge', 
    instance_count       = 1,                 
    max_run              = 2*24*60*60,        
    base_job_name        = job_name,         
    role                 = role,            
    volume_size          = 1024,             
    transformers_version = '4.36.0',        
    pytorch_version      = '2.1.0',          
    py_version           = 'py310',          
    disable_output_compression = True, 
    distribution={"torch_distributed": {"enabled": True}},  
    environment  = {
        "HUGGINGFACE_HUB_CACHE": "/tmp/.cache", 
        "HF_TOKEN": HfFolder.get_token(),     
        "ACCELERATE_USE_FSDP": "1",      
        "FSDP_CPU_RAM_EFFICIENT_LOADING": "
        "FSDP_AUTO_WRAP_POLICY": "TRANSFORMER_BASED_WRAP",
        "FSDP_BACKWARD_PREFETCH": "BACKWARD_PRE",
        "FSDP_STATE_DICT_TYPE": "FULL_STATE_DICT",
        "NCCL_TIMEOUT": "3600",  
        "NCCL_DEBUG": "WARN",   
        "NCCL_IB_TIMEOUT": "3600",
        "NCCL_SOCKET_TIMEOUT": "3600",
        "NCCL_ASYNC_ERROR_HANDLING": "1",
        "NCCL_P2P_LEVEL": "NVL",
        "CUDA_DEVICE_MAX_CONNECTIONS": "1",        
        "MAX_JOBS": "1",                           
        "PYTORCH_CUDA_ALLOC_CONF": "max_split_size_mb:512",
        "TORCH_DISTRIBUTED_DEBUG": "DETAIL",    
    },

    checkpoint_s3_uri=f's3://{bucket}/checkpoints'
)

I'm using a simple training script that leverages the SFTTrainer object, and I'm passing in some basic training parameters.

output_dir: "/opt/ml/model"
model_name_or_path: "Qwen/Qwen2.5-7B-Instruct"
torch_dtype: "bfloat16"
learning_rate: 2e-4
attn_implementation: "flash_attention_2"
packing: false
bf16: "auto"
lr_scheduler_type: "cosine"
warmup_ratio: 0.1
per_device_train_batch_size: 2
gradient_accumulation_steps: 2
max_seq_length: 1024
logging_steps: 10
num_train_epochs: 3
max_steps: 10000
save_steps: 20
load_best_model_at_end: true

Finally, here's how I'm deploying the model after training.

from sagemaker.huggingface import HuggingFaceModel
from sagemaker.huggingface import get_huggingface_llm_image_uri

llm_image = get_huggingface_llm_image_uri(
  "huggingface",
  version="1.4.0",
  session=sess,
)


instance_type = "ml.p4d.24xlarge"
number_of_gpu = 8
health_check_timeout = 1800

config = {
  'HF_MODEL_ID': "/opt/ml/model", 
  'SM_NUM_GPUS': json.dumps(number_of_gpu), 
  'MAX_INPUT_LENGTH': json.dumps(1024), 
  'MAX_TOTAL_TOKENS': json.dumps(2048), 
}

llm_model = HuggingFaceModel(
  role=role,
  image_uri=llm_image,
  model_data='s3:/{bucket}/model.tar.gz',
  env=config
)

This is proving a little tough to debug because the nature of this setup is such that it's hard to figure out exactly what versions of the major libraries are used. The error message suggests the error is coming from https://github.com/huggingface/text-generation-inference/blob/main/server/text_generation_server/models/__init__.py, but it's flagging line 375. In the current version of this file, it looks like this exception is thrown from like 1291.

By printing out llm_image, I can see that I'm using the following Docker image: 763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-tgi-inference:2.1.1-tgi1.4.0-gpu-py310-cu121-ubuntu20.04

Stack trace from SageMaker endpoint

#033[2m2024-11-26T19:35:36.328004Z#033[0m #033[32m INFO#033[0m #033[2mtext_generation_launcher#033[0m#033[2m:#033[0m Args { model_id: "/opt/ml/model", revision: None, validation_workers: 2, sharded: None, num_shard: Some(8), quantize: None, speculate: None, dtype: None, trust_remote_code: false, max_concurrent_requests: 128, max_best_of: 2, max_stop_sequences: 4, max_top_n_tokens: 5, max_input_length: 1024, max_total_tokens: 2048, waiting_served_ratio: 1.2, max_batch_prefill_tokens: 4096, max_batch_total_tokens: None, max_waiting_tokens: 20, hostname: "container-0.local", port: 8080, shard_uds_path: "/tmp/text-generation-server", master_addr: "localhost", master_port: 29500, huggingface_hub_cache: Some("/tmp"), weights_cache_override: None, disable_custom_kernels: false, cuda_memory_fraction: 1.0, rope_scaling: None, rope_factor: None, json_output: false, otlp_endpoint: None, cors_allow_origin: [], watermark_gamma: None, watermark_delta: None, ngrok: false, ngrok_authtoken: None, ngrok_edge: None, tokenizer_config_path: None, env: false }
#033[2m2024-11-26T19:35:36.328028Z#033[0m #033[32m INFO#033[0m #033[2mtext_generation_launcher#033[0m#033[2m:#033[0m Sharding model on 8 processes
#033[2m2024-11-26T19:35:36.328122Z#033[0m #033[32m INFO#033[0m #033[1mdownload#033[0m: #033[2mtext_generation_launcher#033[0m#033[2m:#033[0m Starting download process.
#033[2m2024-11-26T19:35:41.676772Z#033[0m #033[32m INFO#033[0m #033[2mtext_generation_launcher#033[0m#033[2m:#033[0m Files are already present on the host. Skipping download.
#033[2m2024-11-26T19:35:42.336459Z#033[0m #033[32m INFO#033[0m #033[1mdownload#033[0m: #033[2mtext_generation_launcher#033[0m#033[2m:#033[0m Successfully downloaded weights.
#033[2m2024-11-26T19:35:42.336914Z#033[0m #033[32m INFO#033[0m #033[1mshard-manager#033[0m: #033[2mtext_generation_launcher#033[0m#033[2m:#033[0m Starting shard #033[2m#033[3mrank#033[0m#033[2m=#033[0m0#033[0m
#033[2m2024-11-26T19:35:42.337011Z#033[0m #033[32m INFO#033[0m #033[1mshard-manager#033[0m: #033[2mtext_generation_launcher#033[0m#033[2m:#033[0m Starting shard #033[2m#033[3mrank#033[0m#033[2m=#033[0m1#033[0m
#033[2m2024-11-26T19:35:42.337582Z#033[0m #033[32m INFO#033[0m #033[1mshard-manager#033[0m: #033[2mtext_generation_launcher#033[0m#033[2m:#033[0m Starting shard #033[2m#033[3mrank#033[0m#033[2m=#033[0m5#033[0m
#033[2m2024-11-26T19:35:42.337597Z#033[0m #033[32m INFO#033[0m #033[1mshard-manager#033[0m: #033[2mtext_generation_launcher#033[0m#033[2m:#033[0m Starting shard #033[2m#033[3mrank#033[0m#033[2m=#033[0m4#033[0m
#033[2m2024-11-26T19:35:42.337598Z#033[0m #033[32m INFO#033[0m #033[1mshard-manager#033[0m: #033[2mtext_generation_launcher#033[0m#033[2m:#033[0m Starting shard #033[2m#033[3mrank#033[0m#033[2m=#033[0m2#033[0m
#033[2m2024-11-26T19:35:42.337673Z#033[0m #033[32m INFO#033[0m #033[1mshard-manager#033[0m: #033[2mtext_generation_launcher#033[0m#033[2m:#033[0m Starting shard #033[2m#033[3mrank#033[0m#033[2m=#033[0m3#033[0m
#033[2m2024-11-26T19:35:42.337674Z#033[0m #033[32m INFO#033[0m #033[1mshard-manager#033[0m: #033[2mtext_generation_launcher#033[0m#033[2m:#033[0m Starting shard #033[2m#033[3mrank#033[0m#033[2m=#033[0m6#033[0m
#033[2m2024-11-26T19:35:42.337678Z#033[0m #033[32m INFO#033[0m #033[1mshard-manager#033[0m: #033[2mtext_generation_launcher#033[0m#033[2m:#033[0m Starting shard #033[2m#033[3mrank#033[0m#033[2m=#033[0m7#033[0m
#033[2m2024-11-26T19:35:49.859935Z#033[0m #033[31mERROR#033[0m #033[2mtext_generation_launcher#033[0m#033[2m:#033[0m Error when initializing model
Traceback (most recent call last):
  File "/opt/conda/bin/text-generation-server", line 8, in <module>
    sys.exit(app())
  File "/opt/conda/lib/python3.10/site-packages/typer/main.py", line 311, in __call__
    return get_command(self)(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/click/core.py", line 1157, in __call__
    return self.main(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/typer/core.py", line 778, in main
    return _main(
  File "/opt/conda/lib/python3.10/site-packages/typer/core.py", line 216, in _main
    rv = self.invoke(ctx)
  File "/opt/conda/lib/python3.10/site-packages/click/core.py", line 1688, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/opt/conda/lib/python3.10/site-packages/click/core.py", line 1434, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/opt/conda/lib/python3.10/site-packages/click/core.py", line 783, in invoke
    return __callback(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/typer/main.py", line 683, in wrapper
    return callback(**use_params)  # type: ignore
  File "/opt/conda/lib/python3.10/site-packages/text_generation_server/cli.py", line 89, in serve
    server.serve(
  File "/opt/conda/lib/python3.10/site-packages/text_generation_server/server.py", line 235, in serve
    asyncio.run(
  File "/opt/conda/lib/python3.10/asyncio/runners.py", line 44, in run
    return loop.run_until_complete(main)
  File "/opt/conda/lib/python3.10/asyncio/base_events.py", line 636, in run_until_complete
    self.run_forever()
  File "/opt/conda/lib/python3.10/asyncio/base_events.py", line 603, in run_forever
    self._run_once()
  File "/opt/conda/lib/python3.10/asyncio/base_events.py", line 1909, in _run_once
    handle._run()
  File "/opt/conda/lib/python3.10/asyncio/events.py", line 80, in _run
    self._context.run(self._callback, *self._args)
> File "/opt/conda/lib/python3.10/site-packages/text_generation_server/server.py", line 196, in serve_inner
    model = get_model(
  File "/opt/conda/lib/python3.10/site-packages/text_generation_server/models/__init__.py", line 375, in get_model
    raise NotImplementedError("sharded is not supported for AutoModel")
NotImplementedError: sharded is not supported for AutoModel

Expected behavior

The expected behavior is for the SageMaker endpoint to provision successfully.

Related issues

There are a few issues floating around that appear related.

Source: huggingface/text-generation-inference

View original on GitHub View discussion on GitHub