{
  "source_file": "0930",
  "source_platform": "个人博客",
  "author": "Philipp Schmid",
  "title_listed": "SageMaker AWQ + Medusa",
  "keywords": "AWQ,Medusa,加速",
  "domain": "大模型算法",
  "content_type": "教程",
  "is_paid_flag": "否",
  "url": "https://www.philschmid.de/sagemaker-awq-medusa",
  "publish_date": "2024-04",
  "corpus_stage": "推理部署",
  "notes": "AWQ量化+Medusa投机解码加速",
  "id": "ec55d6bdd5e8",
  "collected_at": "2026-10-02T14:29:23",
  "title": "Accelerate Mixtral 8x7B with Speculative Decoding and Quantization on Amazon SageMaker",
  "body_markdown": "Accelerate Mixtral 8x7B with Speculative Decoding and Quantization on Amazon SageMaker\nApril 2, 2024\n9\nminute read\nView Code\nIn this blog you will learn how to accelerate\nmistralai/Mixtral-8x7B-Instruct-v0.1\non Amazon SageMaker using Speculative Decoding (Medusa) and Quantization (AWQ). We are going to use the Hugging Face LLM DLC a purpose-built Inference Container powered by\nText Generation Inference (TGI)\na scalelable, optimized solution for deploying and serving Large Language Models (LLMs). We will prepare the artifacts locally, upload them to Amazon S3, and deploy the model to Amazon SageMaker.\nCombining Medusa and AWQ allows not only to reduce the memory footprint but also to accelerate the inference of Mixtral 8x7B. With this we can deploy Mixtral 8x7B on a single\ng5.12xlarge\ninstance with only 4x NVIDIA A10G GPUs achieving a latency of ~50-60ms per token.\nMixtral 8x7B\nis an open LLM from\nMistral AI\nMixtral-8x7B is a Sparse Mixture of Experts. If you want to learn more about MoEs check out\nMixture of Experts Explained\n. In the blog will cover how to:\nSetup development environment\nRetrieve the new Hugging Face LLM DLC\nPrepare Medusa and AWQ artifacts\nDeploy Mixtral 8x7B to Amazon SageMaker\nRun inference and chat with the model\nLets get started!\nBefore we go into the details, let's first understand what Speculative Decoding and Quantization are and how they can help accelerate Mixtral 8x7B.\nSpeculative Decoding (Medusa)\nMedusa\nis a speculative decoding like technique that can be used to accelerate inference for large language models. It adds extra \"heads\" to LLMs to predict multiple future tokens simultaneously. To use Medusa, you need to train medusa heads on a small dataset preferably similar to the original training data. Medusa can deliver up to a 2x speed increase across a range of LLMs. We are going to use\ntext-generation-inference/Mixtral-8x7B-Instruct-v0.1-medusa\n. The performance improvements depends on the input prompt and can vary.\nQuantization (AWQ)\nAWQ or Activation-aware Weight Quantization\nis a method designed for efficiently compressing and accelerating LLMs by quantizing weights based on the significance of their activations. Resulting in significant reduced memory requirements and improved model latency. AWQ can reduce the memory footprint of a model by up to 4x. We are going to use\nybelkada/Mixtral-8x7B-Instruct-v0.1-AWQ\na awq quantized version of Mixtral 8x7B Instruct.\n1. Setup development environment\nWe are going to use the\nsagemaker\npython SDK to deploy Mixtral to Amazon SageMaker. We need to make sure to have an AWS account configured and the\nsagemaker\npython SDK installed.\nPython\n!\npip install\n\"sagemaker>=2.199.0\"\ntransformers\n--\nupgrade\n--\nquiet\nIf you are going to use Sagemaker in a local environment. You need access to an IAM Role with the required permissions for Sagemaker. You can find\nhere\nmore about it.\nPython\nimport\nsagemaker\nimport\nboto3\nsess = sagemaker.\nSession\n()\n# sagemaker session bucket -> used for uploading data, models and logs\n# sagemaker will automatically create this bucket if it not exists\nsagemaker_session_bucket=\nNone\nif\nsagemaker_session_bucket\nis\nNone\nand\nsess\nis\nnot\nNone\n:\n# set to default bucket if a bucket name is not given\nsagemaker_session_bucket = sess.\ndefault_bucket\n()\ntry\n:\nrole = sagemaker.\nget_execution_role\n()\nexcept\nValueError:\niam = boto3.\nclient\n(\n'iam'\n)\nrole = iam.\nget_role\n(\nRoleName\n=\n'sagemaker_execution_role'\n)[\n'Role'\n][\n'Arn'\n]\nsess = sagemaker.\nSession\n(\ndefault_bucket\n=sagemaker_session_bucket)\nprint(\nf\n\"sagemaker role arn:\n{\nrole\n}\n\"\n)\nprint(\nf\n\"sagemaker session region:\n{\nsess.boto_region_name\n}\n\"\n)\n2. Retrieve the new Hugging Face LLM DLC\nCompared to deploying regular Hugging Face models we first need to retrieve the container uri and provide it to our\nHuggingFaceModel\nmodel class with a\nimage_uri\npointing to the image. To retrieve the new Hugging Face LLM DLC in Amazon SageMaker, we can use the\nget_huggingface_llm_image_uri\nmethod provided by the\nsagemaker\nSDK. This method allows us to retrieve the URI for the desired Hugging Face LLM DLC based on the specified\nbackend\n,\nsession\n,\nregion\n, and\nversion\n. You can find the available versions\nhere\nPython\nfrom\nsagemaker.huggingface\nimport\nget_huggingface_llm_image_uri\n# retrieve the llm image uri\nllm_image =\nget_huggingface_llm_image_uri\n(\n\"huggingface\"\n,\nversion\n=\n\"1.4.2\"\n)\n# print ecr image uri\nprint(\nf\n\"llm image uri:\n{\nllm_image\n}\n\"\n)\n3. Prepare Medusa and AWQ artifacts\nIn this step we are going to prepare the Medusa and AWQ artifacts. Hugging Face Text Generation Inference container supports\nMedusa\nand AWQ. For Medusa we need to provide the a\nconfig.json\nwith the medusa specifications, including\nbase_model_name_or_path\nand\nmedusa_lm_head.safetensors\n. If we only wanted to deploy the Medusa model we could use\ntext-generation-inference/Mixtral-8x7B-Instruct-v0.1-medusa\nas\nHF_MODEL_ID\n.\nBut since we want to combine Medusa and AWQ and deploy from Amazon S3 we need to modify the\nconfig.json\nto include the AWQ model. We are going to use the\nybelkada/Mixtral-8x7B-Instruct-v0.1-AWQ\nmodel.\nThis will lead to a directory structure like this:\nmodel/\n├── llm/\n│   ├── config.json\n│   ├── model-00001-of-00005.safetensors\n│   ├── ...\n├── config.json\n├── medusa_lm_head.safetensors\n└── ...\nWe will first load the Medusa model and AWQ model from Hugging Face, modify the\nconfig.json\nand then upload the artifacts to Amazon S3.\nPython\nimport\njson\nimport\nos\nfrom\nhuggingface_hub\nimport\nsnapshot_download\nfrom\nsagemaker.s3\nimport\nS3Uploader\ntmp_dir =\n\"./tmp\"\nmedusa_repository =\n\"text-generation-inference/Mixtral-8x7B-Instruct-v0.1-medusa\"\n# Medusa model\n# https://huggingface.co/TheBloke/Mixtral-8x7B-Instruct-v0.1-AWQ/discussions/6\nllm_repository =\n\"ybelkada/Mixtral-8x7B-Instruct-v0.1-AWQ\"\n# AWQ LLM model\nsnapshot_download\n(\nrepo_id\n=medusa_repository,\nlocal_dir\n=tmp_dir)\nsnapshot_download\n(\nrepo_id\n=llm_repository,\nlocal_dir\n=os.path.\njoin\n(tmp_dir,\n\"llm\"\n),\nignore_patterns\n=\n\"*.bin\"\n)\n# rewrite meudsa base model value\nwith\nopen(os.path.\njoin\n(tmp_dir,\n\"config.json\"\n),\n\"r\"\n)\nas\nf:\ndata = json.\nload\n(f)\ndata[\n\"base_model_name_or_path\"\n] =\n\"/opt/ml/model/llm/\"\n# path to llm model in side Amazon SageMaker\nwith\nopen(os.path.\njoin\n(tmp_dir,\n\"config.json\"\n),\n\"w\"\n)\nas\nf_out:\njson.\ndump\n(data,\nindent\n=\n2\n,\nfp\n=f_out)\n# upload the model to s3\ns3_path = S3Uploader.\nupload\n(\nlocal_path\n=tmp_dir,\ndesired_s3_uri\n=\nf\n\"s3://\n{\nsess.\ndefault_bucket\n()\n}\n/medusa/mixtral\"\n)\n4. Deploy Mixtral 8x7B to Amazon SageMaker\nAfter we have prepared the Medusa and AWQ artifacts we can deploy Mixtral 8x7B to Amazon SageMaker. We will use the\nHuggingFaceModel\nmodel class and define our endpoint configuration including the\nhf_model_id\n,\ninstance_type\netc. We will use a\ng5.12xlarge\ninstance type, which has 4 NVIDIA A10G GPUs and 96GB of GPU memory.\nWe make sure to have\nQUANTIZE\nset to\nawq\nand that\nHF_MODEL_ID\nis set to\n/opt/ml/model\n, since SageMaker will load the model from S3 to there.\nPython\nimport\njson\nfrom\nsagemaker.huggingface\nimport\nHuggingFaceModel\n# sagemaker config\ninstance_type =\n\"ml.g5.12xlarge\"\nnumber_of_gpu =\n4\nhealth_check_timeout =\n300\n# Define Model and Endpoint configuration parameter\nconfig = {\n'HF_MODEL_ID'\n:\n\"/opt/ml/model\"\n,\n# path to where sagemaker stores the model\n'SM_NUM_GPUS'\n: json.\ndumps\n(number_of_gpu),\n# Number of GPU used per replica\n'MAX_INPUT_LENGTH'\n: json.\ndumps\n(\n8000\n),\n# Max length of input text\n'MAX_BATCH_PREFILL_TOKENS'\n: json.\ndumps\n(\n16384\n),\n# Number of tokens for the prefill operation.\n'MAX_TOTAL_TOKENS'\n: json.\ndumps\n(\n16384\n),\n# Max length of the generation (including input text)\n'QUANTIZE'\n:\n\"awq\"\n# Quantization method\n}\n# create HuggingFaceModel with the image uri\nllm_model =\nHuggingFaceModel\n(\nrole\n=role,\n# path to s3 bucket with model, we are not using a compressed model\nmodel_data\n={\n'S3DataSource'\n:{\n'S3Uri'\n: s3_path +\n\"/\"\n,\n'S3DataType'\n:\n'S3Prefix'\n,\n'CompressionType'\n:\n'None'\n}},\nimage_uri\n=llm_image,\nenv\n=config\n)\nAfter we have created the\nHuggingFaceModel\nwe can deploy it to Amazon SageMaker using the\ndeploy\nmethod. TGI will automatically distribute and shard the model across all GPUs.\nPython\n# Deploy model to an endpoint\n# https://sagemaker.readthedocs.io/en/stable/api/inference/model.html#sagemaker.model.Model.deploy\nllm = llm_model.\ndeploy\n(\ninitial_instance_count\n=\n1\n,\ninstance_type\n=instance_type,\ncontainer_startup_health_check_timeout\n=health_check_timeout,\n# 10 minutes to be able to load the model\n)\nSageMaker will now create our endpoint and deploy the model to it. This can takes a 10-15 minutes.\n5. Run inference and chat with the model\nAfter our endpoint is deployed we can run inference on it. We will use the\npredict\nmethod from the\npredictor\nto run inference on our endpoint. We can inference with different parameters to impact the generation. Parameters can be defined as in the\nparameters\nattribute of the payload. You can find supported parameters in the\nhere\nor in the open api specification of the TGI in the\nswagger documentation\nThe\nmistralai/Mixtral-8x7B-Instruct-v0.1\nis a conversational chat model meaning we can chat with it using the following prompt:\n<s> [INST] User Instruction 1 [/INST] Model answer 1</s> [INST] User instruction 2 [/INST]\nLets see, if Mixtral can come up with some cool ideas for the summer.\nPython\nfrom\ntransformers\nimport\nAutoTokenizer\ntok = AutoTokenizer.\nfrom_pretrained\n(\n\"ybelkada/Mixtral-8x7B-Instruct-v0.1-AWQ\"\n)\n# Prompt to generate\nmessages = [\n{\n\"role\"\n:\n\"user\"\n,\n\"content\"\n:\n\"How can i make a good american cheese cake? Explain it step-by-step and include time estimates.\"\n},\n]\n# Generation arguments\npayload = {\n\"do_sample\"\n:\nTrue\n,\n\"top_p\"\n:\n0.6\n,\n\"temperature\"\n:\n0.9\n,\n\"top_k\"\n:\n50\n,\n\"max_new_tokens\"\n:\n1024\n,\n\"repetition_penalty\"\n:\n1.03\n,\n\"return_full_text\"\n:\nFalse\n,\n\"stop\"\n: [\n\"</s>\"\n]\n}\nOkay lets test it.\nPython\nchat = llm.\npredict\n({\n\"inputs\"\n:tok.\napply_chat_template\n(messages,\ntokenize\n=\nFalse\n,\nadd_generation_prompt\n=\nTrue\n),\n# convert messages to model input\n\"parameters\"\n:payload\n})\nprint(chat[\n0\n][\n\"generated_text\"\n])\nAwesome, We have successfully deployed Mixtral 8x7B to Amazon SageMaker with Medusa and AWQ on a\ng5.12xlarge\ninstance. Thats not possible with the regular Mixtral 8x7B model. If we now look up the latency of the model on Cloudwatch we can see that we have a latency of ~50-60ms per token.\n6. Clean up\nTo clean up, we can delete the model and endpoint.\nPython\nllm.\ndelete_model\n()\nllm.\ndelete_endpoint\n()\nConclusion\nLeveraging Medusa and AWQ allowed us to deploy Mixtral-8x7B on a\ng5.12xlarge\n, reducing costs by ~3x (from\ng5.48xlarge\n) and improving latency to 50-60ms per token. This is a great example of how you can go beyond the standard deployment of Hugging Face models on Amazon SageMaker and leverage the power of the Hugging Face LLM DLC to accelerate inference and reduce costs.\nThanks for reading! If you have any questions, feel free to contact me on\nTwitter\nor\nLinkedIn\n.",
  "metadata": {
    "title": "Accelerate Mixtral 8x7B with Speculative Decoding and Quantization on Amazon SageMaker",
    "description": "In this blog post you will learn how to accelerate Mixtral using Speculative Decoding (Medusa) and Quantization (AWQ).",
    "keywords": null,
    "author": "Philipp Schmid",
    "og:title": "Accelerate Mixtral 8x7B with Speculative Decoding and Quantization on Amazon SageMaker",
    "og:description": "In this blog post you will learn how to accelerate Mixtral using Speculative Decoding (Medusa) and Quantization (AWQ).",
    "og:url": "https://www.philschmid.de/sagemaker-awq-medusa",
    "og:image": "https://www.philschmid.de/static/blog/sagemaker-awq-medusa/thumbnail.jpg",
    "og:image:width": "1200",
    "og:image:height": "630",
    "og:image:alt": "Accelerate Mixtral 8x7B with Speculative Decoding and Quantization on Amazon SageMaker",
    "og:type": "article",
    "twitter:card": "summary_large_image",
    "twitter:title": "Accelerate Mixtral 8x7B with Speculative Decoding and Quantization on Amazon SageMaker",
    "twitter:description": "In this blog post you will learn how to accelerate Mixtral using Speculative Decoding (Medusa) and Quantization (AWQ).",
    "twitter:image": "https://www.philschmid.de/static/blog/sagemaker-awq-medusa/thumbnail.jpg",
    "article:published_time": "2024-04-02"
  },
  "paywall": {
    "listed_paid": false,
    "detected": false,
    "marker": "",
    "public_only": false,
    "policy": "付费/会员内容只保留页面公开可见部分，不绕过付费墙"
  },
  "crawl": {
    "success": true,
    "status_code": 200,
    "error": ""
  },
  "code_blocks": [
    {
      "file": "attachments/ec55d6bdd5e8/code/body-000.txt",
      "chars": 64,
      "language": "text"
    },
    {
      "file": "attachments/ec55d6bdd5e8/code/body-001.txt",
      "chars": 734,
      "language": "text"
    },
    {
      "file": "attachments/ec55d6bdd5e8/code/body-002.txt",
      "chars": 235,
      "language": "text"
    },
    {
      "file": "attachments/ec55d6bdd5e8/code/body-003.txt",
      "chars": 143,
      "language": "text"
    },
    {
      "file": "attachments/ec55d6bdd5e8/code/body-004.txt",
      "chars": 1020,
      "language": "text"
    },
    {
      "file": "attachments/ec55d6bdd5e8/code/body-005.txt",
      "chars": 989,
      "language": "text"
    },
    {
      "file": "attachments/ec55d6bdd5e8/code/body-006.txt",
      "chars": 320,
      "language": "text"
    },
    {
      "file": "attachments/ec55d6bdd5e8/code/body-007.txt",
      "chars": 90,
      "language": "text"
    },
    {
      "file": "attachments/ec55d6bdd5e8/code/body-008.txt",
      "chars": 524,
      "language": "text"
    },
    {
      "file": "attachments/ec55d6bdd5e8/code/body-009.txt",
      "chars": 203,
      "language": "text"
    },
    {
      "file": "attachments/ec55d6bdd5e8/code/body-010.txt",
      "chars": 40,
      "language": "text"
    }
  ],
  "images": [],
  "audio": [],
  "videos": [],
  "links": [
    {
      "href": "https://github.com/huggingface/notebooks/blob/main/sagemaker/31_speculative_decoding_quantization_inference/sagemaker-notebook.ipynb",
      "label": "View Code",
      "type": "a"
    },
    {
      "href": "https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1",
      "label": "mistralai/Mixtral-8x7B-Instruct-v0.1",
      "type": "a"
    },
    {
      "href": "https://github.com/huggingface/text-generation-inference",
      "label": "Text Generation Inference (TGI)",
      "type": "a"
    },
    {
      "href": "https://huggingface.co/mistralai",
      "label": "Mistral AI",
      "type": "a"
    },
    {
      "href": "https://huggingface.co/blog/moe",
      "label": "Mixture of Experts Explained",
      "type": "a"
    },
    {
      "href": "https://github.com/FasterDecoding/Medusa/tree/main",
      "label": "Medusa",
      "type": "a"
    },
    {
      "href": "https://huggingface.co/text-generation-inference/Mixtral-8x7B-Instruct-v0.1-medusa",
      "label": "text-generation-inference/Mixtral-8x7B-Instruct-v0.1-medusa",
      "type": "a"
    },
    {
      "href": "https://github.com/mit-han-lab/llm-awq",
      "label": "AWQ or Activation-aware Weight Quantization",
      "type": "a"
    },
    {
      "href": "https://huggingface.co/ybelkada/Mixtral-8x7B-Instruct-v0.1-AWQ",
      "label": "ybelkada/Mixtral-8x7B-Instruct-v0.1-AWQ",
      "type": "a"
    },
    {
      "href": "https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-roles.html",
      "label": "here",
      "type": "a"
    },
    {
      "href": "https://github.com/aws/deep-learning-containers/blob/master/available_images.md",
      "label": "here",
      "type": "a"
    },
    {
      "href": "https://huggingface.co/docs/text-generation-inference/conceptual/speculation",
      "label": "Medusa",
      "type": "a"
    },
    {
      "href": "https://huggingface.co/text-generation-inference/Mixtral-8x7B-Instruct-v0.1-medusa/blob/main/config.json",
      "label": "config.json",
      "type": "a"
    },
    {
      "href": "https://www.philschmid.de/sagemaker-llama-llm",
      "label": "here",
      "type": "a"
    },
    {
      "href": "https://huggingface.github.io/text-generation-inference/",
      "label": "swagger documentation",
      "type": "a"
    },
    {
      "href": "https://twitter.com/_philschmid",
      "label": "Twitter",
      "type": "a"
    },
    {
      "href": "https://www.linkedin.com/in/philipp-schmid-a6a2bb196/",
      "label": "LinkedIn",
      "type": "a"
    }
  ],
  "repair_version": 2,
  "repaired_at": "2026-10-14T10:52:54Z",
  "source_url": "https://www.philschmid.de/sagemaker-awq-medusa",
  "papers": [],
  "actual_video_files": [],
  "body_chars": 10973,
  "expected_body_images": 0,
  "discovery_status": "ok",
  "delivery_status": "ok",
  "missing_reason": "",
  "raw_html": "attachments/ec55d6bdd5e8/pages/source.html",
  "local_html": "attachments/ec55d6bdd5e8/pages/article.html",
  "body_base_path": "attachments/ec55d6bdd5e8/pages/",
  "body_markdown_path": "attachments/ec55d6bdd5e8/pages/article.md"
}