AWS Study Group

Tối ưu hóa Mixtral 8x7B trên Amazon SageMaker với AWS Inferentia2

Tác giả: Lior Sadan & Stenio de Lima Ferreira | Ngày: 15/4/2025

Danh mục: Machine Learning, Amazon SageMaker, AWS Inferentia, Hướng dẫn kỹ thuật

Các tổ chức đang tìm kiếm cách triển khai các mô hình ngôn ngữ lớn (LLM) mạnh mẽ như Mixtral 8x7B một cách hiệu quả về chi phí và hiệu suất. Bài viết này sẽ hướng dẫn chi tiết cách tối ưu hóa và triển khai Mixtral 8x7B trên AWS Inferentia2 với Amazon SageMaker.

Tại sao chọn AWS Inferentia2 cho Mixtral 8x7B?

Ưu điểm của kiến trúc Mixture-of-Experts (MoE)

Mixtral 8x7B sử dụng kiến trúc MoE với 8 chuyên gia (experts), cho phép:

  • Hiệu suất cao với 46.7 tỷ tham số
  • Tối ưu hóa tài nguyên thông qua expert parallelism
  • Khả năng mở rộng trên nhiều NeuronCore

AWS Inferentia2: Chip AI tối ưu

Thông số kỹ thuật:

  • Mỗi NeuronCore: 16GB High Bandwidth Memory (HBM)
  • Phiên bản inf2.24xlarge: 12 NeuronCore (6 chip Inferentia2)
  • Tối ưu cho inference với độ trễ thấp và thông lượng cao

Yêu cầu hệ thống và chuẩn bị

Bước 1: Thiết lập Hugging Face Access

Tạo tài khoản và token

  1. Đăng ký tài khoản tại Hugging Face
  2. Tạo Access Token với quyền đọc/ghi tại Settings → Access Tokens
  3. Yêu cầu quyền truy cập model Mixtral-8x7B-Instruct-v0.1

Lưu ý bảo mật: Luôn giữ token an toàn và không chia sẻ công khai

Bước 2: Khởi tạo EC2 Instance

Sử dụng CloudFormation Template (Khuyến nghị)

# Clone repository mẫu
git clone https://github.com/aws-samples/sample-optimizing-mixtral-8x7B-on-amazon-sagemaker-with-aws-inferentia2

Hoặc tạo thủ công qua Console

Cấu hình khuyến nghị:

  • Instance Type: inf2.24xlarge (12 NeuronCore)
  • AMI: Hugging Face Neuron Deep Learning AMI
  • Storage: 512GB EBS (cho model lớn)
  • Security Group: Cho phép SSH từ IP của bạn

Kết nối và kiểm tra

# SSH với port forwarding cho Jupyter
ssh -i "<pem-file>" ubuntu@<instance-dns> -L 8888:127.0.0.1:8888

# Kiểm tra NeuronCore
neuron-ls

Output mong đợi cho inf2.24xlarge:

instance-type: inf2.24xlarge
+--------+-------+--------+----------+-------+
| NEURON | CORES | MEMORY | DEVICES  | BDF   |
+--------+-------+--------+----------+-------+
| 0-5    | 2     | 32 GB  | Various  | ...   |
+--------+-------+--------+----------+-------+

Tính toán yêu cầu bộ nhớ

Công thức tính toán

Total Memory = bytes_per_parameter × number_of_parameters

Cho Mixtral 8x7B:

  • Tham số: 46.7 tỷ
  • Với FP16: 93.4GB cho weights
  • KV Cache: ~0.5GB (batch=1, seq_len=1024)
  • Tổng: ~94GB

Tensor Parallelism

Yêu cầu: Phân phối trên 8 NeuronCore

  • Mỗi NeuronCore: 16GB HBM
  • Tensor parallel levels hỗ trợ: 8, 16, 32
  • Lựa chọn tối ưu: 8 cores cho inf2.24xlarge

Biên dịch Mixtral 8x7B

Bước 1: Khởi tạo Container

# Chạy Neuron TGI Container
docker run -it --entrypoint /bin/bash \
  --net=host -v $(pwd):$(pwd) -w $(pwd) \
  --device=/dev/neuron0 \
  --device=/dev/neuron1 \
  --device=/dev/neuron2 \
  --device=/dev/neuron3 \
  --device=/dev/neuron4 \
  --device=/dev/neuron5 \
  ghcr.io/huggingface/neuronx-tgi:0.0.25

Bước 2: Đăng nhập Hugging Face

# Trong container
huggingface-cli login --token hf_<your-token>

Bước 3: Biên dịch Model với Optimum-CLI

optimum-cli export neuron \
  --model mistralai/Mixtral-8x7B-Instruct-v0.1 \
  --batch_size 1 \
  --sequence_length 1024 \
  --auto_cast_type fp16 \
  --num_cores 8 \
  ./neuron_model_path

Giải thích tham số

Tham sốGiá trịÝ nghĩa
batch_size1Số chuỗi input đồng thời
sequence_length1024Độ dài token tối đa
auto_cast_typefp16Kiểu dữ liệu (tối ưu memory)
num_cores8Số NeuronCore sử dụng

Thời gian biên dịch: 10-20 phút

Bước 4: Upload Model đã biên dịch

# Upload lên Hugging Face Hub
huggingface-cli upload <user_id>/Mixtral-8x7B-Instruct-v0.1 ./neuron_model_path ./

# Hoặc lưu trên S3
aws s3 cp ./neuron_model_path s3://your-bucket/mixtral-compiled/ --recursive

Triển khai trên Amazon SageMaker

Thiết lập IAM Permissions

Tạo IAM Role cho EC2

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Principal": {
        "Service": [
          "ec2.amazonaws.com",
          "sagemaker.amazonaws.com"
        ]
      },
      "Action": "sts:AssumeRole"
    }
  ]
}

Policies cần thiết:

  • AmazonSageMakerFullAccess
  • IAMReadOnlyAccess

Security Best Practice: Sử dụng least privilege principle trong production

Khởi chạy Jupyter Notebook

# Cài đặt Jupyter trong container
pip install ipykernel jupyter notebook environment_kernels
python3 -m ipykernel install --user --name aws_neuron_venv_pytorch

# Khởi chạy server
jupyter notebook

Truy cập: http://localhost:8888 qua SSH tunnel

Code triển khai SageMaker

1. Import thư viện và cấu hình

import os
import sagemaker
from sagemaker.huggingface import get_huggingface_llm_image_uri, HuggingFaceModel

# Cấu hình
os.environ['AWS_DEFAULT_REGION'] = 'us-east-1'
sess = sagemaker.Session()
role = sagemaker.get_execution_role()

# Lấy container image
llm_image = get_huggingface_llm_image_uri("huggingface-neuronx", version="0.0.25")

2. Cấu hình Model

# Cấu hình model parameters
config = {
    "HF_MODEL_ID": "user_id/Mixtral-8x7B-Instruct-v0.1",
    "HF_NUM_CORES": "8",
    "HF_AUTO_CAST_TYPE": "fp16",
    "MAX_BATCH_SIZE": "1",
    "MAX_INPUT_LENGTH": "1000",
    "MAX_TOTAL_TOKENS": "1024",
    "MESSAGES_API_ENABLED": "true",
    "HUGGING_FACE_HUB_TOKEN": "hf_<your-token>"
}

# Tạo HuggingFaceModel
llm_model = HuggingFaceModel(
    role=role,
    image_uri=llm_image,
    env=config
)

3. Deploy Endpoint

# Cấu hình instance
instance_type = "ml.inf2.24xlarge"
health_check_timeout = 2400  # Thời gian load model
volume_size = 512  # GB

# Deploy model
llm_model._is_compiled_model = True
llm = llm_model.deploy(
    initial_instance_count=1,
    instance_type=instance_type,
    container_startup_health_check_timeout=health_check_timeout,
    volume_size=volume_size
)

4. Test Inference

# Tạo prompt
messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "What is deep learning?"}
]

parameters = {
    "model": "user_id/Mixtral-8x7B-Instruct-v0.1",
    "top_p": 0.6,
    "temperature": 0.9,
    "max_tokens": 1000,
}

# Gửi request
response = llm.predict({"messages": messages, **parameters})
print(response["choices"][0]["message"]["content"].strip())

Tối ưu hóa Performance và Cost

Monitoring và Scaling

# Kiểm tra endpoint status
endpoints = sess.sagemaker_client.list_endpoints()
for endpoint in endpoints['Endpoints']:
    print(f"Endpoint: {endpoint['EndpointName']} - Status: {endpoint['EndpointStatus']}")

Cost Optimization

Chiến lược tiết kiệm:

  1. Auto Scaling: Cấu hình scale down khi không sử dụng
  2. Spot Instances: Sử dụng cho development/testing
  3. Reserved Capacity: Cho production workloads
# Xóa endpoint khi không sử dụng
llm.delete_endpoint()
llm.delete_model()

Troubleshooting và Best Practices

Common Issues

1. Out of Memory Errors

# Kiểm tra memory usage
neuron-monitor

2. Slow Inference

# Tối ưu batch size và sequence length
config["MAX_BATCH_SIZE"] = "4"  # Tăng throughput
config["MAX_INPUT_LENGTH"] = "512"  # Giảm memory usage

3. Model Loading Timeout

# Tăng timeout
health_check_timeout = 3600  # 1 hour

Best Practices

Security

  • ✅ Sử dụng IAM roles thay vì hardcode credentials
  • ✅ Encrypt data in transit và at rest
  • ✅ Sử dụng VPC endpoints cho S3/Hugging Face

Performance

  • ✅ Monitor CloudWatch metrics
  • ✅ Sử dụng appropriate instance types
  • ✅ Implement caching strategies

Cost Management

  • ✅ Set up billing alerts
  • ✅ Use lifecycle policies cho S3 storage
  • ✅ Regular review và cleanup unused resources

Kết luận và Next Steps

Việc triển khai Mixtral 8x7B trên AWS Inferentia2 mang lại:

Lợi ích chính

  • Cost Effective: Giảm 50-70% chi phí so với GPU instances
  • High Performance: Inference tối ưu với độ trễ thấp
  • Scalability: Auto-scaling theo demand
  • Managed Service: SageMaker handle infrastructure

Phát triển tiếp

Explore Advanced Features:

Tài nguyên tham khảo

Documentation

Sample Code

Community

Về tác giả:

Lior Sadan – Senior Solution Architect tại AWS, chuyên về AI/ML infrastructure và storage solutions.

Stenio de Lima Ferreira – Senior Solution Architect với 15+ năm kinh nghiệm cloud infrastructure, DevOps và data science.

←
→

Leave a comment