Cách Amazon Search M5 tiết kiệm 30% chi phí đào tạo LLM bằng cách sử dụng AWS Trainium | Dịch vụ web của Amazon

Trong nhiều thập kỷ, Amazon đã đi tiên phong và đổi mới máy học (ML), mang lại những trải nghiệm thú vị cho khách hàng của mình. Ngay từ những ngày đầu tiên, Amazon đã sử dụng ML cho nhiều trường hợp sử dụng khác nhau như đề xuất sách, tìm kiếm và phát hiện gian lận. Tương tự như phần còn lại của ngành, những tiến bộ của phần cứng được tăng tốc đã cho phép các nhóm của Amazon theo đuổi kiến ​​trúc mô hình bằng cách sử dụng mạng thần kinh và học sâu (DL).

Chương trình M5 trong Amazon Search sở hữu chiến lược học tập khám phá cho Amazon và xây dựng các mô hình quy mô lớn trên đa ngôn ngữ, đa ngôn ngữ, đa thực thể, đa nhiệm và đa phương thức như văn bản, hình ảnh và video. Chương trình M5 đã và đang cung cấp các mô hình nền tảng quy mô lớn và nhúng phổ quát cho hàng trăm nhóm ML trên Amazon trong khi vẫn duy trì các biện pháp kiểm soát chặt chẽ về tối ưu hóa chi phí. Để đạt được điều này, nhóm M5 thường xuyên đánh giá các kỹ thuật mới để giảm chi phí.

Giống như nhiều tổ chức ML, máy gia tốc phần lớn được sử dụng để tăng tốc quá trình đào tạo và suy luận DL. Khi AWS ra mắt các chương trình tăng tốc được xây dựng có mục đích với bản phát hành đầu tiên của Suy luận AWS vào năm 2020, nhóm M5 nhanh chóng bắt đầu sử dụng chúng để triển khai khối lượng công việc sản xuất hiệu quả hơn, tiết kiệm cả chi phí và giảm độ trễ. Năm ngoái, AWS đã ra mắt Đào tạo AWS bộ tăng tốc giúp tối ưu hóa hiệu suất trên mỗi chi phí để phát triển và xây dựng các mô hình DL thế hệ tiếp theo. Trong bài đăng này, chúng tôi thảo luận về cách M5 có thể giảm 30% chi phí đào tạo mô hình của họ và chia sẻ một số phương pháp hay nhất mà chúng tôi đã học được trong quá trình thực hiện.

Phiên bản Trainium

Với những tiến bộ về các máy gia tốc được xây dựng có mục đích, Amazon cũng cung cấp các máy gia tốc hấp dẫn dưới dạng AWS Inferentia và Trainium. Đúng như tên gọi của chúng, những con chip này được tối ưu hóa để đáp ứng nhu cầu tương ứng của khối lượng công việc suy luận và đào tạo. Để đào tạo quy mô lớn các mô hình nền tảng đạt kích thước hàng tỷ tham số, Trainium Phiên bản Trn1 và Trn1n là sự lựa chọn lý tưởng do đặc điểm của chúng. Các phiên bản Trn1 được hỗ trợ bởi công nghệ tiên tiến nhất NeuronCore-v2và có rất nhiều bộ nhớ và khả năng tính toán tăng tốc. Các phiên bản Trn1n cũng có thể được chọn để có lượng băng thông mạng lớn hơn (1,600 Gbs), do đó rất phù hợp để đào tạo hiệu suất với mục đích tối ưu hóa chi phí.

Để sử dụng bộ tăng tốc, bạn cần có một lớp phần mềm để hỗ trợ chúng. Với chip Trn và Inf, SDK thần kinh AWS mở khóa các máy gia tốc được xây dựng có mục đích của Amazon với sự trợ giúp của PyTorch XLA. PyTorch XLA chuyển đổi chế độ háo hức của PyTorch sang triển khai dựa trên biểu đồ ở chế độ lười biếng. Những biểu đồ này sau đó được sử dụng và biên dịch thêm để sử dụng với máy gia tốc. PyTorch Neuron (một phần của Neuron SDK) cho phép người dùng PyTorch đào tạo mô hình của họ trên Trainium NeuronCores bằng một vài dòng mã.

Mô hình và khối lượng công việc

Nhóm M5 đào tạo và triển khai các mô hình nền tảng cũng như cách trình bày phổ quát để hỗ trợ các nhóm khác nhau trên Amazon trong việc mang lại niềm vui cho Amazon.com khách hàng. Một mô hình như vậy là mô hình bộ mã hóa văn bản, theo sau là perceptron nhiều lớp (MLP) với các tương tác tính năng rõ ràng hoặc tiềm ẩn được xác định bởi kiến ​​trúc mạng thần kinh với hàng trăm triệu tham số có thể huấn luyện. Mô hình này được đào tạo trên hàng tỷ mã thông báo và được sử dụng để tạo ra hàng triệu lượt nhúng trong cài đặt suy luận hàng loạt ngoại tuyến. Các phần nhúng này là đầu vào cho dịch vụ Amazon cấp 1 hướng tới khách hàng.

Cơ sở hạ tầng cho việc sử dụng đường ống sản xuất Hàng loạt AWS với chiến lược xếp hàng chia sẻ công bằng, sử dụng cụm nhiều nút trn1.32xlarge có hỗ trợ EFA làm công cụ tính toán cho việc đào tạo mô hình. Về mặt chức năng, quy trình sản xuất thực hiện đào tạo mô hình gia tăng, đánh giá mô hình được đào tạo và suy luận hàng loạt ngoại tuyến trên mô hình được đào tạo, tất cả đều sử dụng PyTorch làm thư viện DL cơ bản.

Các mục tiêu

Làm hài lòng khách hàng của chúng tôi là một nguyên lý quan trọng nhất. Do tính chất hướng tới khách hàng của quy trình, điều quan trọng là phải đáp ứng tất cả các thỏa thuận cấp độ dịch vụ (SLA) mà không bị thoái lui. Chúng tôi đã xác định hai tiêu chí chấp nhận quan trọng để điều chỉnh quy trình sản xuất GPU hiện có của mình và chuyển nó sang Trainium:

• Chất lượng mô hình – Chất lượng mẫu mã của chúng tôi tác động trực tiếp đến trải nghiệm của khách hàng. Chúng tôi yêu cầu phải có chênh lệch ít hơn 0.1% về chất lượng mô hình giữa GPU và Trainium.
• Thông lượng đào tạo – Chúng tôi đào tạo lặp đi lặp lại các mô hình của mình theo định kỳ để cung cấp trải nghiệm mới nhất cho khách hàng. Chúng tôi yêu cầu phải đạt được sự hội tụ mô hình trong khoảng thời gian xác định trước (chẳng hạn như 1 tuần) để đáp ứng SLA sản xuất của chúng tôi.

Trong các phần sau, chúng tôi chia sẻ hành trình làm việc ngược lại từ tiêu chí này cũng như những bài học của chúng tôi để hỗ trợ khối lượng công việc sản xuất trên quy mô Amazon.

Kịch bản đào tạo

Trước khi bắt đầu đào tạo mô hình, chúng ta cần thực hiện các thay đổi đối với tập lệnh đào tạo để làm cho nó tuân thủ XLA. Với kích thước của mô hình, chúng tôi sử dụng song song dữ liệu phân tán (DDP) để huấn luyện mô hình. DDP cho phép chúng tôi tăng hiệu suất đào tạo mô hình bằng cách tăng quy mô số lượng máy được sử dụng để chạy đào tạo mô hình mà không cần bất kỳ thay đổi mã nào. Chúng tôi đã làm theo hướng dẫn được cung cấp trong Hướng dẫn đào tạo Neuron PyTorch MLP để thêm các cấu trúc dành riêng cho XLA vào tập lệnh đào tạo của chúng tôi. Những thay đổi mã này rất dễ thực hiện. Sau đây là một số bài học kỹ thuật quan trọng từ bài tập đã cải thiện đáng kể thông lượng mô hình của chúng tôi:

• Vị trí của xm.mark_step() – xm.mark_step() biên dịch và chạy các biểu đồ tính toán được thu thập một cách lười biếng. Kêu gọi mark_step quá nhiều lần sẽ dẫn đến số lượng lớn hơn các biểu đồ nhỏ, trong khi gọi nó quá ít lần sẽ dẫn đến một số biểu đồ nhỏ nhưng lớn. Tùy thuộc vào ứng dụng của bạn, thông lượng và việc triển khai đào tạo mô hình của bạn sẽ khác nhau tùy theo vị trí của bạn xm.mark_step(). Việc triển khai của chúng tôi đặt một xm.mark_step() sau một lượt tiến và lùi và một lượt sau bước tối ưu hóa.
• Gói tải dữ liệu bằng trình tải thiết bị đa xử lý XLA - Đây là bước quan trọng rất dễ bị bỏ qua. Trình tải thiết bị đa xử lý torch_xla.distributed.parallel_loader.MpDeviceLoader tải dữ liệu đào tạo trên mỗi thiết bị XLA với các tùy chọn tải trước và tải dữ liệu chồng chéo với các lần chạy thiết bị để cải thiện thông lượng. Trình tải thiết bị cũng gọi xm.mark_step() và do đó có thể xây dựng biểu đồ để tải dữ liệu từ máy chủ vào thiết bị.

Biên soạn cho Trainium

Theo truyền thống, chu trình phát triển mô hình với GPU bao gồm việc thực hiện các thay đổi đối với mô hình hoặc tập lệnh huấn luyện và chạy trực tiếp nó trên thiết bị GPU. Các máy gia tốc như Trainium sử dụng XLA yêu cầu một bước bổ sung trước khi có thể chạy đào tạo mô hình trên máy gia tốc. Đồ thị tính toán XLA chỉ có thể chạy sau khi chúng được biên dịch. Nói chung, có hai cách để thực hiện quá trình biên dịch này: Trước thời gian (AOT), trong đó bạn theo dõi và biên dịch tất cả các biểu đồ trước rồi chạy chúng hoặc Just In Time (JIT), trong đó các biểu đồ được theo dõi, biên dịch và chạy khi chúng đang gặp phải. SDK Neuron cung cấp cả hai thứ này. Thông thường, quá trình biên dịch AOT được thực hiện trước tiên. Đồ thị sau đó được chạy sau quá trình biên dịch này. Nếu gặp biểu đồ mới, thời gian chạy Neuron sẽ gọi trình biên dịch JIT trước khi chạy chúng. Để thực hiện biên dịch AOT, SDK Neuron cung cấp nơron_parallel_compile, một tiện ích biên dịch trích xuất các biểu đồ từ lần chạy thử tập lệnh đào tạo và thực hiện biên dịch AOT song song.

Một khía cạnh quan trọng của quá trình biên dịch AOT là đảm bảo rằng không có biểu đồ tính toán mới nào được tạo trong quá trình đào tạo. Một nguồn của các biểu đồ tính toán mới (và do đó được biên dịch lại) là các hình dạng động của các đợt huấn luyện trong quá trình huấn luyện mô hình. Chúng tôi nhận thấy rằng việc sử dụng các hình dạng tĩnh và các lô có kích thước cố định sẽ loại bỏ việc tổng hợp thời gian đào tạo và cải thiện đáng kể thông lượng đào tạo mà không ảnh hưởng đến độ chính xác của mô hình. Bằng cách thực thi các ràng buộc như vậy trong quá trình đào tạo, chúng tôi nhận thấy rằng chỉ cần 4–5 bước đào tạo mô hình, một bước xác thực mô hình và kiểm tra điểm mô hình một lần để theo dõi tất cả các biểu đồ trong quá trình biên dịch AOT. Điều quan trọng cần lưu ý là SDK Neuron không ngừng phát triển và trong tương lai cũng sẽ hỗ trợ các hình dạng động.

Hơn nữa, các biểu đồ được biên dịch được lưu trữ trong Bộ nhớ đệm liên tục của Neuron trên đĩa hoặc trong một Dịch vụ lưu trữ đơn giản của Amazon (Amazon S3). Điều này đặc biệt hữu ích cho khối lượng công việc sản xuất mà kiến ​​trúc mô hình và cấu hình đào tạo không thay đổi. Do đó, chi phí biên dịch chỉ phát sinh một lần. Sử dụng bộ đệm cũng đơn giản như đặt cờ môi trường:

export NEURON_COMPILE_CACHE_URL="s3://BUCKET/KEY"

Trình biên dịch Neuron cũng cung cấp ba tùy chọn tối ưu hóa cấp độ trình biên dịch (O1, O2, O3) để cân bằng thời gian biên dịch và thông lượng chạy mô hình. O1 cho phép tối ưu hóa cốt lõi trên biểu đồ điện toán và giảm thiểu thời gian biên dịch, O3 cung cấp thông lượng chạy mô hình được cải thiện với chi phí thời gian biên dịch cao hơn và O2 (tùy chọn mặc định) là sự cân bằng giữa cả hai. Đối với trường hợp sử dụng của chúng tôi, chúng tôi đã sử dụng tính năng tối ưu hóa O1 và nhận thấy thời gian biên dịch giảm 86% mà không thay đổi các chỉ số về độ chính xác của mô hình, đồng thời quan sát thấy thông lượng giảm khoảng 5–7% so với tối ưu hóa mặc định (O2). Tùy thuộc vào trường hợp sử dụng, bạn có thể chọn mức độ tối ưu hóa khác nhau.

Để tóm tắt, chúng tôi đã sử dụng các cờ sau để biên dịch:

NEURON_CC_FLAGS="--target trn1 --auto-cast all --auto-cast-type bf16 --model-type transformer --optlevel O1"

Khả năng tương thích điểm kiểm tra

Khi quá trình biên dịch hoàn tất thành công, chúng ta có thể tiến hành đào tạo các mô hình của mình trên Trainium. Như đã đề cập trước đó, chúng tôi huấn luyện dần dần các mô hình của mình, nghĩa là chúng tôi tải điểm kiểm tra mô hình đã huấn luyện trước đó và tiếp tục huấn luyện với dữ liệu mới. PyTorch và PyTorch XLA cho phép chuyển đổi liền mạch giữa các máy gia tốc thông qua khả năng tương tác của điểm kiểm tra. Việc di chuyển linh hoạt giữa GPU và Trainium cho phép chúng tôi tải mô hình GPU trước đó và đào tạo trên máy Trainium một cách liền mạch. Điều này rất quan trọng để đảm bảo rằng chúng tôi có thể khởi tạo mô hình của mình với mô hình được đào tạo tốt nhất trước đó mà không có bất kỳ thời gian ngừng sản xuất nào hoặc mất đi độ chính xác của mô hình.

Vì mô hình GPU đã được lưu bằng các tiện ích lưu mô hình PyTorch tiêu chuẩn nên chúng tôi có thể sử dụng tiện ích tải điểm kiểm tra PyTorch để tải mô hình GPU trên các thiết bị Trainium.

Ví dụ: trên GPU/CPU, bạn có thể lưu mô hình bằng mã sau:

torch.save(model.state_dict(), PATH)

Sau đó bạn tải lại mô hình trên Trainium:

import torch_xla.core.xla_model as xm
xla_device = xm.xla_device()
model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(xla_device)

Tương tự, bạn có thể lưu mô hình trên Trainium bằng đoạn mã sau:

import torch_xla.core.xla_model as xm
# automatically moves the data to CPU for the master device
xm.save(model.state_dict(), PATH) 

Và tải lại mô hình trên GPU/CPU:

model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device) # can be any device

Trên thực tế, vì chúng tôi sử dụng DDP để huấn luyện mô hình nên việc tải mô hình không phụ thuộc vào số lượng máy được sử dụng để huấn luyện điểm kiểm tra trước đó. Điều này cho phép chúng tôi mở rộng quy mô nhóm Trn1 theo chiều ngang mà không thay đổi mã hoặc ảnh hưởng bất lợi đến việc đào tạo mô hình. Các điểm kiểm tra dựa trên PyTorch này có thể được sử dụng trực tiếp hoặc thậm chí được viết kịch bản theo đuốc cho các trường hợp sử dụng suy luận trên AWS Inferentia2 hoặc các trình tăng tốc khác.

Hoạt động ổn định

Không thể nhấn mạnh đủ rằng việc chạy khối lượng công việc trong sản xuất đòi hỏi phải đáp ứng nhiều SLA. Đối với trường hợp sử dụng của chúng tôi, ngoài chất lượng mô hình và thông lượng đào tạo SLA, điều bắt buộc là quy trình sản xuất phải hoạt động ổn định, nghĩa là thời gian ngừng hoạt động và gián đoạn ở mức tối thiểu trong quá trình đào tạo, đánh giá và suy luận mô hình.

Giống như quy trình dựa trên GPU hiện có, chúng tôi đã thêm nhiều cơ chế để giúp quy trình hoạt động ổn định. Trước khi bắt đầu đào tạo mô hình, chúng tôi chạy nhiều bài kiểm tra độ chính xác để đánh giá tình trạng của máy. Các thử nghiệm này thường bao gồm các thao tác tensor đơn giản để xác minh tình trạng hoạt động của các thiết bị máy gia tốc. Chúng tôi nhận thấy rằng đối với hoạt động đào tạo phân tán, điều quan trọng là phải chạy thử nghiệm để xác minh hoạt động giao tiếp chung giữa các phiên bản. Chúng tôi đã sử dụng Bộ thử nghiệm NCCOM từ SDK Neuron để đạt được điều này, chạy nhiều hoạt động khác nhau như tập hợp toàn bộ, giảm toàn bộ và giảm phân tán.

Ngay cả sau khi làm theo những đề xuất mà chúng tôi đã đề cập, chúng tôi nhận thấy rằng các vấn đề nhất thời là không thể tránh khỏi trong bất kỳ quy trình nào, bất kể máy gia tốc cơ bản là gì. Để xây dựng khả năng phục hồi trong bất kỳ quy trình đào tạo nào, chúng tôi khuyên bạn nên xây dựng cơ chế thử lại để giải quyết các vấn đề tiềm ẩn này. Chúng tôi sử dụng Thử lại tự động hàng loạt AWS để thử lại các công việc gặp phải lỗi tạm thời trong quá trình đào tạo mô hình. Những lần khởi động lại này có thể tốn kém nếu gặp phải lỗi vào cuối quá trình đào tạo. Để giải quyết vấn đề này, chúng tôi đã điều chỉnh tập lệnh đào tạo của mình để tải điểm kiểm tra mô hình đã đào tạo trước đó và tiếp tục đào tạo từ thời điểm đó. Với chức năng này, chúng tôi có thể tích cực khởi động lại các công việc đào tạo không thành công với chi phí tối thiểu.

Với các cơ chế phục hồi này được áp dụng, chúng tôi có thể đạt được tỷ lệ thành công 98.5% cho khối lượng công việc của mình trên Trn1, tương đương với tỷ lệ thành công của quy trình GPU hiện tại của chúng tôi.

Kết quả

Để xác thực tính chính xác của các mô hình, chúng tôi đã khởi tạo hai mô hình từ cùng một điểm kiểm tra GPU và huấn luyện một mô hình trên Trainium và mô hình còn lại trên GPU tương đương. Cả hai mô hình đều được huấn luyện với các siêu tham số huấn luyện giống nhau. Tập dữ liệu được sử dụng để tính toán số liệu là tập dữ liệu loại trừ và chúng tôi đánh giá độ chính xác của mô hình trên tập dữ liệu này sau mỗi N bước toàn cầu. Trục X là bước tổng thể và trục Y là độ chính xác của mô hình. Chúng tôi quan sát thấy sự khác biệt ít hơn 0.1% về độ chính xác của mô hình tại mỗi điểm trong biểu đồ sau.

Hơn nữa, để đánh giá hiệu quả chi phí của việc đào tạo mô hình, chúng tôi muốn so sánh thời gian đồng hồ treo tường thực hiện để đạt được sự hội tụ của mô hình. Chúng tôi tin rằng điều này mang lại cái nhìn thực tế hơn về việc tiết kiệm chi phí so với các biện pháp như giá mỗi token, FLOPS/đô la đạt được và các yếu tố khác. Xét thời gian huấn luyện của trn1.32xl và so sánh Đám mây điện toán đàn hồi Amazon (Amazon EC2), chúng tôi nhận thấy rằng Trainium cung cấp chi phí rẻ hơn tới 30% cho mô hình hội tụ.

Kết luận

Có nhiều yếu tố cần xem xét khi đánh giá các chương trình tăng tốc khác nhau cho khối lượng công việc DL của bạn. Một số điều quan trọng nhất là chất lượng mô hình, thông lượng, chi phí và tính sẵn có. Điều tối quan trọng là phải đảm bảo rằng chất lượng và thông lượng mô hình của bạn không bị ảnh hưởng do công cụ tăng tốc bạn chọn.

Nhờ sự hợp tác và hợp tác của chúng tôi với nhóm Annapurna Neuron, nhóm Amazon Search M5 đã có thể tiết kiệm tới 30% chi phí bằng cách chuyển sang Trainium. Nhóm có thể sử dụng Trainium và đạt được chất lượng mô hình cũng như thông lượng tương đương với các máy gia tốc tương đương trên thị trường. Khả năng tương tác của điểm kiểm tra và thay đổi mã tối thiểu với sự hỗ trợ cho XLA đã cho phép M5 lựa chọn giữa nhiều máy gia tốc cho khối lượng công việc của họ. Điều này đã cho phép nhóm M5 tận dụng sức mạnh tính toán lớn của Trainium và xây dựng các giải pháp tăng tốc bất khả tri để làm hài lòng khách hàng của Amazon.com. Từ quan điểm vận hành, Trainium đã được chứng minh là có khả năng hỗ trợ các dịch vụ cấp 1 ở quy mô Amazon. Nhóm M5 tiếp tục chuyển nhiều khối lượng công việc hơn sang Trainium để cung cấp các mô hình tốt nhất cho Amazon với chi phí thấp nhất.

Tóm lại, nhóm M5 đã có thể thực hiện chương trình đào tạo ML cấp sản xuất, tiết kiệm chi phí bằng cách thêm Trainium vào nhóm máy gia tốc. Chúng tôi khuyến khích bạn xem qua Trainium và các thiết bị Neuron khác như AWS Inferentia để tận dụng lợi ích của silicon Amazon được xây dựng có mục đích cho khối lượng công việc ML. Bắt đầu dễ dàng với một trong nhiều hướng dẫn có các mô hình khác nhau, như Llama 2, có sẵn trên Trainium.

Về các tác giả

Abhinandan Patni là Kỹ sư phần mềm cao cấp tại Amazon Search. Ông tập trung vào việc xây dựng hệ thống và công cụ để đào tạo học sâu phân tán có thể mở rộng và suy luận theo thời gian thực.

Công viên James là Kiến trúc sư giải pháp tại Amazon Web Services. Anh làm việc với Amazon.com để thiết kế, xây dựng và triển khai các giải pháp công nghệ trên AWS và đặc biệt quan tâm đến AI và học máy. Trong thời gian rảnh rỗi, anh ấy thích tìm kiếm những nền văn hóa mới, những trải nghiệm mới và luôn cập nhật những xu hướng công nghệ mới nhất. Bạn có thể tìm thấy anh ấy trên LinkedIn.

Jerry Mannil là kỹ sư phần mềm tại Amazon Search. Ông nỗ lực cải thiện tính hiệu quả, mạnh mẽ và khả năng mở rộng của cơ sở hạ tầng đào tạo phân tán.

Ken Su là kỹ sư phần mềm tại Amazon Search. Anh ấy nỗ lực cải thiện hiệu quả đào tạo và quy trình đào tạo phân tán có thể mở rộng. Ngoài công việc, anh ấy thích đi bộ đường dài và chơi tennis.

RJ là một Kỹ sư của Amazon. Anh ấy xây dựng và tối ưu hóa hệ thống cho các hệ thống phân tán dành cho đào tạo và nỗ lực tối ưu hóa các hệ thống áp dụng để giảm độ trễ cho Suy luận ML. Ngoài công việc, anh ấy đang khám phá việc sử dụng Generative AI để xây dựng các công thức nấu ăn.

• Phân phối nội dung và PR được hỗ trợ bởi SEO. Được khuếch đại ngay hôm nay.
• PlatoData.Network Vertical Generative Ai. Trao quyền cho chính mình. Truy cập Tại đây.
• PlatoAiStream. Thông minh Web3. Kiến thức khuếch đại. Truy cập Tại đây.
• Trung tâmESG. Than đá, công nghệ sạch, Năng lượng, Môi trường Hệ mặt trời, Quản lý chất thải. Truy cập Tại đây.
• PlatoSức khỏe. Tình báo thử nghiệm lâm sàng và công nghệ sinh học. Truy cập Tại đây.
• nguồn: https://aws.amazon.com/blogs/machine-learning/how-amazon-search-m5-saved-30-for-llm-training-cost-by-using-aws-trainium/