Học Sâu Phân Tán (Distributed Deep Learning)

Huấn luyện phân tán chia nhỏ tính toán ra nhiều GPU và nhiều máy để huấn luyện các mô hình quá lớn hoặc quá chậm đối với một thiết bị đơn lẻ. File này bao quát độ chính xác hỗn hợp, song song dữ liệu, song song mô hình, song song đường ống, ZeRO, FSDP, song song tensor, và các nguyên thủy giao tiếp như all-reduce — những thứ thiết yếu để huấn luyện LLM quy mô lớn.

• Huấn luyện một mạng nơ-ron lớn trên một GPU đơn lẻ cuối cùng sẽ đụng tường. Mô hình có thể không vừa bộ nhớ, hoặc việc huấn luyện có thể mất hàng tháng. Huấn luyện phân tán chia công việc ra nhiều thiết bị (GPU, TPU, hoặc cả những cỗ máy hoàn chỉnh) để huấn luyện nhanh hơn và huấn luyện được các mô hình lớn hơn. File này bao quát các kỹ thuật giúp điều đó khả thi.

• Để hiểu tại sao phân tán lại quan trọng, hãy bắt đầu từ chi phí tính toán của việc huấn luyện. Một lượt truyền xuôi đơn lẻ qua một tầng dày (dense layer) với \(d_{\text{in}}\) đầu vào và \(d_{\text{out}}\) đầu ra trên một batch gồm \(B\) mẫu đòi hỏi xấp xỉ \(2 \cdot B \cdot d_{\text{in}} \cdot d_{\text{out}}\) FLOPs (phép toán dấu phẩy động): một phép nhân và một phép cộng cho mỗi phần tử của ma trận đầu ra. Lượt truyền ngược tốn xấp xỉ gấp đôi lượt truyền xuôi (tính gradient đối với cả đầu vào và trọng số), nên một bước huấn luyện trên tầng dày vào khoảng \(6 \cdot B \cdot d_{\text{in}} \cdot d_{\text{out}}\) FLOPs.

• Với một tầng Transformer có số chiều ẩn \(d\), khối tự-chú ý (self-attention) gồm bốn phép chiếu (Q, K, V và đầu ra), mỗi cái tốn \(O(B \cdot n \cdot d^2)\) FLOPs (với \(n\) là độ dài chuỗi), cộng thêm phép tính ma trận chú ý ở \(O(B \cdot n^2 \cdot d)\). Khối truyền thẳng (feed-forward) có hai tầng dày, thường mở rộng lên \(4d\) rồi thu về: \(O(B \cdot n \cdot 8d^2)\). Tổng mỗi tầng: xấp xỉ \(O(B \cdot n \cdot 12d^2 + B \cdot n^2 \cdot d)\). Nhân với số lớp và bạn sẽ thấy vì sao huấn luyện các mô hình quy mô GPT tiêu tốn hàng nghìn GPU-giờ.

• Bức tường bộ nhớ thường là ràng buộc chặt chẽ hơn. Trong lúc huấn luyện, bộ nhớ GPU phải chứa đồng thời bốn thứ:

• Tham số (Parameters): trọng số của mô hình. Một mô hình 7 tỷ tham số ở dạng FP32 (4 byte mỗi tham số) cần tới 28 GB chỉ riêng cho trọng số.
• Gradient: cùng kích thước với tham số. Thêm 28 GB nữa.
• Trạng thái bộ tối ưu hóa (Optimizer states): Adam duy trì hai bộ đệm bổ sung (ước lượng moment bậc một và bậc hai), mỗi bộ bằng kích thước của tham số. Chúng được giữ ở FP32 để ổn định số học, ngay cả khi mô hình dùng độ chính xác thấp hơn. Với mô hình 7B của ta, con số này là \(2 \times 28 = 56\) GB.

• Bản kích hoạt (Activations): các giá trị trung gian được lưu trong lượt truyền xuôi để dùng cho lượt truyền ngược. Kích thước phụ thuộc vào kích thước batch, độ dài chuỗi và độ rộng mô hình. Đây thường là thành phần lớn nhất và tăng tuyến tính theo kích thước batch.

• Với mô hình 7B dùng Adam FP32: 28 (tham số) + 28 (gradient) + 56 (bộ tối ưu hóa) = 112 GB, trước cả khi tính đến bản kích hoạt. Một GPU A100 80 GB đơn lẻ không thể chứa nổi. Đó là lý do các chiến lược phân tán là thiết yếu.

• Huấn luyện độ chính xác hỗn hợp (Mixed precision training) là lớp phòng thủ đầu tiên. Thay vì lưu mọi thứ ở FP32 (số chấm động 32-bit), bạn huấn luyện dùng FP16 hoặc BF16 (16-bit) cho các lượt truyền xuôi và truyền ngược, trong khi giữ một bản sao trọng số chính (master copy) ở FP32 cho bước cập nhật bộ tối ưu hóa.

• FP16 có độ chính xác cao (phần định trị 10-bit) nhưng dải giá trị hạn chế, có thể gây tràn/ngoài dải (overflow/underflow). Việc nhân tỷ lệ mất mát (loss scaling — nhân mất mát với một hệ số lớn trước lượt truyền ngược, rồi chia gradient cho cùng hệ số đó) giúp giảm thiểu điều này.

• BF16 (brain float) có cùng dải số mũ với FP32 (số mũ 8-bit) nhưng độ chính xác thấp hơn (phần định trị 7-bit). Nó hầu như không bao giờ bị tràn và ít khi cần nhân tỷ lệ mất mát, giúp việc sử dụng đơn giản hơn. BF16 là mặc định cho huấn luyện Transformer hiện đại.

• Độ chính xác hỗn hợp xấp xỉ halve (giảm một nửa) bộ nhớ cho bản kích hoạt và gradient (những chi phí chiếm ưu thế trong các lượt truyền xuôi/ngược), đồng thời giữ trạng thái bộ tối ưu hóa ở FP32 để ổn định số học.

• Song song dữ liệu (Data parallelism) là chiến lược phân tán đơn giản nhất. Bạn nhân bản toàn bộ mô hình lên \(N\) GPU, chia mỗi mini-batch thành \(N\) phần bằng nhau, và gửi một phần cho mỗi GPU. Mỗi GPU chạy lượt truyền xuôi và truyền ngược trên phần của nó một cách độc lập. Sau đó các gradient được trung bình hóa trên tất cả GPU (dùng thao tác all-reduce), và mỗi GPU cập nhật bản sao cục bộ của mô hình.

• Xét từ góc độ mô hình, việc này tương đương với huấn luyện với một mini-batch lớn gấp \(N\) lần. Nếu mỗi GPU xử lý một batch kích thước \(B\), thì kích thước batch hiệu dụng là \(N \cdot B\).

• Việc trung bình hóa gradient có thể thực hiện đồng bộ hoặc bất đồng bộ. SGD đồng bộ (Synchronous SGD) chờ tất cả GPU hoàn thành trước khi trung bình hóa, đảm bảo tương đương số học với huấn luyện đơn GPU trên một batch lớn hơn. Nhược điểm là GPU chậm nhất (kẻ "trễ nhịp" - "straggler") kéo lùi tất cả.

• SGD bất đồng bộ (Asynchronous SGD) cho phép mỗi GPU cập nhật một máy chủ tham số chia sẻ độc lập, mà không cần chờ. Cách này loại bỏ vấn đề trễ nhịp nhưng lại sinh ra "gradient lỗi thời" (stale gradients): một GPU có thể tính gradient dựa trên các tham số hơi cũ. Gradient lỗi thời thêm nhiễu và có thể làm chậm sự hội tụ. Trong thực tế, SGD đồng bộ với giao tiếp hiệu quả là lựa chọn ưu tiên.

• Tích lũy gradient (Gradient accumulation) là một thủ thuật phần mềm để mô phỏng các kích thước batch lớn hơn trên phần cứng hạn chế. Thay vì cập nhật một lần trên mỗi mini-batch, bạn chạy vài lượt truyền xuôi/truyền ngược và cộng dồn các gradient, rồi mới thực hiện một lần cập nhật. Cách này cho kết quả giống hệt một batch lớn mà không cần thêm bộ nhớ GPU cho bản kích hoạt (chỉ một mini-batch bản kích hoạt nằm trong bộ nhớ tại một thời điểm).

• Khi bản thân mô hình quá lớn để vừa một GPU đơn lẻ, bạn cần song song mô hình (model parallelism). Có hai biến thể chính.

• Song song tensor (Tensor parallelism) chia nhỏ các lớp riêng lẻ ra nhiều GPU. Một phép nhân ma trận lớn \(Y = XW\) có thể được chia theo cột: phân hoạch \(W\) thành \([W_1, W_2]\) trên hai GPU, tính \(Y_1 = XW_1\) và \(Y_2 = XW_2\) song song, rồi nối tiếp. Cách này hoạt động cho các phép chiếu chú ý và các tầng truyền thẳng. Nó đòi hỏi giao tiếp nhanh giữa các GPU (thường là NVLink trong một node) vì các kết quả trung gian phải được kết hợp tại mỗi lớp.

• Song song đường ống (Pipeline parallelism) gán các lớp khác nhau cho các GPU khác nhau. GPU 0 chạy các lớp 1-4, GPU 1 chạy các lớp 5-8, và cứ thế. Dữ liệu chảy qua đường ống như một dây chuyền lắp ráp. Cách tiếp cận ngây thơ sinh ra một "bong bóng đường ống" (pipeline bubble): trong khi GPU 0 xử lý lượt truyền xuôi cho micro-batch 1, các GPU 1-3 ngồi không. Việc chia nhỏ thành các micro-batch (Micro-batching) giảm thiểu điều này bằng cách chia mini-batch thành các micro-batch nhỏ hơn chảy qua đường ống theo thứ tự, giữ tất cả GPU bận rộn phần lớn thời gian.

• Song song lai (Hybrid parallelism) kết hợp song song dữ liệu, tensor và đường ống. Một thiết lập mô hình lớn điển hình có thể dùng song song tensor bên trong một node (8 GPU nối bởi NVLink nhanh), song song đường ống xuyên các node, và song song dữ liệu xuyên các nhóm node. Đó là cách các mô hình như GPT-4 và Llama được huấn luyện.

• Hiệu quả của huấn luyện phân tán phụ thuộc rất nhiều vào giao tiếp (communication). Thao tác then chốt là all-reduce: với một giá trị trên mỗi \(N\) GPU, tính tổng (hoặc trung bình) và phân phối kết quả cho tất cả GPU.

• Một all-reduce ngây thơ gửi tất cả dữ liệu về một GPU, cộng tổng lại, rồi phát ngược trở lại. Cách này có độ phức tạp giao tiếp \(O(N)\) và tạo ra điểm thắt nút ở root.

• Ring all-reduce hiệu quả hơn nhiều. Sắp xếp \(N\) GPU thành một vòng. Mỗi GPU chia dữ liệu của nó thành \(N\) phần. Trong \(N - 1\) bước, mỗi GPU gửi một phần cho hàng xóm và nhận một phần từ hàng xóm còn lại, cộng dồn các tổng bộ phận. Sau thêm \(N - 1\) bước nữa, tổng đầy đủ được truyền đến tất cả GPU. Tổng dữ liệu truyền trên mỗi GPU: \(2(N-1)/N\) lần kích thước dữ liệu, tiến gần đến \(2\times\) khi \(N\) tăng. Điều quan trọng là nó không tăng theo \(N\), giúp tối ưu băng thông.

• Máy chủ tham số (Parameter servers) là một kiến trúc thay thế, trong đó các node máy chủ chuyên dụng nắm giữ các tham số mô hình. Các worker tính gradient và gửi cho máy chủ, máy chủ cập nhật tham số rồi gửi lại. Cách này đơn giản nhưng có thể tạo ra điểm thắt nút giao tiếp tại máy chủ.

• NCCL (NVIDIA Collective Communications Library) là thư viện chuẩn cho giao tiếp GPU-to-GPU. Nó cung cấp các cài đặt tối ưu của all-reduce, all-gather, broadcast và các thao tác tập thể khác, tự động chọn thuật toán tốt nhất cho tô-pô mạng.

• Định luật tỷ lệ (Scaling laws) mô tả cách hiệu năng mô hình cải thiện cùng tính toán, dữ liệu và kích thước mô hình. Định luật tỷ lệ ban đầu của Kaplan et al. (2020) thấy rằng mất mát giảm theo hàm lũy thừa với mỗi đại lượng:

\[L(N) \propto N^{-\alpha_N}, \quad L(D) \propto D^{-\alpha_D}, \quad L(C) \propto C^{-\alpha_C}\]

• với \(N\) là số tham số, \(D\) là kích thước tập dữ liệu, và \(C\) là ngân sách tính toán.

• Định luật tỷ lệ Chinchilla (Hoffmann et al., 2022) chỉ ra rằng hầu hết các mô hình đều được huấn luyện thiếu: với một ngân sách tính toán cho trước, bạn nên huấn luyện một mô hình nhỏ hơn trên nhiều dữ liệu hơn so với suy nghĩ trước đây. Tỷ lệ tối ưu xấp xỉ 20 token cho mỗi tham số. Một mô hình 7B nên "thấy" khoảng 140B token, chứ không phải 300B token như Llama 1 dùng với mô hình 65B. Phát hiện này đã chuyển hướng lĩnh vực sang huấn luyện "tối ưu tính toán" (compute-optimal).

• Hỗn hợp chuyên gia (Mixture of Experts - MoE) là một kiến trúc mở rộng dung lượng mô hình mà không tỷ lệ theo tính toán. Thay vì một mạng truyền thẳng cho mỗi tầng Transformer, bạn có \(N\) mạng "chuyên gia" (mỗi cái là một FFN chuẩn). Một mạng cổng (gating network - router) xem xét mỗi token và gửi nó đến top-\(K\) chuyên gia (thường \(K = 1\) hoặc \(K = 2\)).

• Tổng số tham số lớn hơn nhiều (vì bạn có \(N\) chuyên gia), nhưng FLOPs trên mỗi token gần như không đổi (vì chỉ \(K\) chuyên gia được kích hoạt mỗi token). Ví dụ, Mixtral 8x7B có 47B tổng tham số nhưng chỉ dùng khoảng 13B mỗi lượt truyền xuôi, cho hiệu năng của một mô hình lớn hơn nhiều với giá của một mô hình nhỏ hơn.

• MoE đặt ra những thách thức. Cân bằng tải (Load balancing): nếu bộ định tuyến gửi hầu hết token đến cùng một chuyên gia, các chuyên gia khác bị lãng phí. Một mất mát phụ trợ khuyến khích định tuyến đồng đều. Giao tiếp: các chuyên gia khác nhau có thể nằm trên các GPU khác nhau, nên việc định tuyến token đòi hỏi giao tiếp all-to-all, vốn tốn kém.

• Dung sai lỗi (Fault tolerance) cực kỳ quan trọng khi quá trình huấn luyện kéo dài hàng tuần hoặc hàng tháng trên hàng nghìn GPU. Nếu một GPU duy nhất hỏng, bạn không muốn mất toàn bộ tiến trình. Tạo checkpoint (Checkpointing) định kỳ lưu trọng số mô hình, trạng thái bộ tối ưu hóa và trạng thái huấn luyện (tốc độ học, số bước, vị trí dữ liệu) xuống đĩa. Nếu có lỗi xảy ra, bạn khởi động lại từ checkpoint gần nhất.

• Tạo checkpoint gradient (Gradient checkpointing - còn gọi là tính toán lại bản kích hoạt) là một tối ưu bộ nhớ, không phải cơ chế dung sai lỗi. Trong lượt truyền xuôi, thay vì lưu mọi bản kích hoạt cho lượt truyền ngược, bạn chỉ lưu các bản kích hoạt tại một số checkpoint nhất định. Trong lượt truyền ngược, bạn tính lại các bản kích hoạt bị thiếu từ các checkpoint. Cách này đổi tính toán lấy bộ nhớ: nó tăng chi phí truyền xuôi khoảng 33% nhưng có thể giảm bộ nhớ bản kích hoạt đi một hệ số \(\sqrt{L}\) (với \(L\) là số lớp).

• Gộp tất cả lại, huấn luyện một mô hình tiên phong kết hợp tất cả các kỹ thuật này: độ chính xác hỗn hợp BF16, song song dữ liệu xuyên hàng nghìn GPU với ring all-reduce, song song tensor trong các node, song song đường ống xuyên các node, tạo checkpoint gradient để giảm bộ nhớ, MoE cho hiệu quả tham số, và tạo checkpoint định kỳ để dung sai lỗi. Kỹ nghệ hệ thống cũng thử thách ngang với thiết kế thuật toán.

• Tóm lại bộ công cụ huấn luyện phân tán:

Kỹ thuật	Tác dụng	Đánh đổi
Độ chính xác hỗn hợp (BF16)	Halve bộ nhớ cho bản kích hoạt/gradient	Sai khác số học nhẹ
Song song dữ liệu	Mở rộng kích thước batch xuyên GPU	Chi phí giao tiếp để đồng bộ gradient
Song song tensor	Chia nhỏ các lớp xuyên GPU	Đòi hỏi kết nối liên thông nhanh
Song song đường ống	Chia nhỏ các giai đoạn mô hình xuyên GPU	Bong bóng đường ống (tính toán lãng phí)
Tích lũy gradient	Mô phỏng các batch lớn	Chậm hơn (nhiều lượt truyền xuôi/ngược)
Tạo checkpoint gradient	Giảm bộ nhớ bản kích hoạt	Tăng ~33% tính toán
Ring all-reduce	Trung bình hóa gradient hiệu quả	Giới hạn băng thông với mô hình lớn
MoE	Nhiều dung lượng hơn, cùng FLOPs	Cân bằng tải, độ phức tạp định tuyến
Định luật tỷ lệ	Định hướng phân bổ tính toán	Thực nghiệm, có thể không đúng mọi quy mô

Bài tập lập trình (dùng CoLab hoặc notebook)

1. Tính FLOPs và yêu cầu bộ nhớ cho một tầng Transformer. Với số chiều ẩn \(d\), độ dài chuỗi \(n\), kích thước batch \(B\), và số lớp, ước tính tổng chi phí huấn luyện.

   import jax.numpy as jnp
   
   def transformer_layer_flops(d, n, B):
       """Xấp xỉ FLOPs cho một lượt truyền xuôi của một tầng Transformer."""
       # Các phép chiếu QKV: 3 * (B * n * d * d) * 2 (nhân-cộng)
       qkv_flops = 3 * 2 * B * n * d * d
       # Chú ý: (B * n * n * d) * 2 cho QK^T, (B * n * n * d) * 2 cho attn*V
       attn_flops = 2 * 2 * B * n * n * d
       # Phép chiếu đầu ra: (B * n * d * d) * 2
       out_flops = 2 * B * n * d * d
       # FFN: hai lớp, d->4d và 4d->d: 2 * (B * n * d * 4d) * 2
       ffn_flops = 2 * 2 * B * n * d * 4 * d
       return qkv_flops + attn_flops + out_flops + ffn_flops
   
   def transformer_layer_memory(d, n, B, dtype_bytes=2):
       """Xấp xỉ bộ nhớ bản kích hoạt (byte) cho một lớp."""
       # QKV: 3 * B * n * d
       qkv_mem = 3 * B * n * d * dtype_bytes
       # Trọng số chú ý: B * heads * n * n (xấp xỉ B * n * n * sizeof)
       attn_mem = B * n * n * dtype_bytes
       # Khớp trung gian FFN: B * n * 4d
       ffn_mem = B * n * 4 * d * dtype_bytes
       return qkv_mem + attn_mem + ffn_mem
   
   # Ví dụ: quy mô GPT-2
   d, n, B, L = 1024, 1024, 8, 24
   fwd_flops = transformer_layer_flops(d, n, B)
   total_flops = 3 * L * fwd_flops  # 3x cho truyền xuôi + truyền ngược
   act_mem = L * transformer_layer_memory(d, n, B)
   param_count = L * (12 * d * d + 13 * d)  # xấp xỉ
   
   print(f"Model: d={d}, n={n}, B={B}, L={L}")
   print(f"Parameters: {param_count / 1e6:.0f}M")
   print(f"FLOPs per step: {total_flops / 1e12:.2f} TFLOPs")
   print(f"Activation memory: {act_mem / 1e9:.2f} GB (BF16)")
   print(f"Parameter memory (FP32): {param_count * 4 / 1e9:.2f} GB")
   print(f"Adam optimizer memory: {param_count * 8 / 1e9:.2f} GB")
   print(f"Total training memory: {(param_count * 16 + act_mem) / 1e9:.2f} GB")
   

2. Mô phỏng huấn luyện song song dữ liệu. Chia một tập dữ liệu ra nhiều "GPU ảo", tính gradient độc lập, trung bình hóa chúng, và xác minh kết quả khớp với huấn luyện đơn GPU.

   import jax
   import jax.numpy as jnp
   
   # Mô hình tuyến tính đơn giản: y = wx + b
   key = jax.random.PRNGKey(0)
   X = jax.random.normal(key, (64, 4))
   w_true = jnp.array([1.0, -2.0, 3.0, 0.5])
   y = X @ w_true + 0.1 * jax.random.normal(key, (64,))
   
   def loss_fn(w, X, y):
       return jnp.mean((X @ w - y) ** 2)
   
   grad_fn = jax.grad(loss_fn)
   
   # Đơn GPU: gradient full batch
   w = jnp.zeros(4)
   grad_single = grad_fn(w, X, y)
   
   # Song song dữ liệu: chia thành 4 "GPU"
   n_gpus = 4
   chunk_size = len(X) // n_gpus
   grads = []
   for i in range(n_gpus):
       X_chunk = X[i*chunk_size:(i+1)*chunk_size]
       y_chunk = y[i*chunk_size:(i+1)*chunk_size]
       grads.append(grad_fn(w, X_chunk, y_chunk))
   
   # All-reduce: trung bình hóa gradient
   grad_parallel = jnp.mean(jnp.stack(grads), axis=0)
   
   print("Single-GPU gradient:", grad_single)
   print("Data-parallel gradient (avg):", grad_parallel)
   print(f"Match: {jnp.allclose(grad_single, grad_parallel, atol=1e-5)}")
   
   # Huấn luyện cả hai và so sánh
   w_single, w_parallel = jnp.zeros(4), jnp.zeros(4)
   lr = 0.1
   for step in range(100):
       w_single = w_single - lr * grad_fn(w_single, X, y)
   
       grads = [grad_fn(w_parallel, X[i*chunk_size:(i+1)*chunk_size],
                        y[i*chunk_size:(i+1)*chunk_size]) for i in range(n_gpus)]
       avg_grad = jnp.mean(jnp.stack(grads), axis=0)
       w_parallel = w_parallel - lr * avg_grad
   
   print(f"\nAfter 100 steps:")
   print(f"Single-GPU weights: {w_single}")
   print(f"Data-parallel weights: {w_parallel}")
   print(f"Max difference: {jnp.max(jnp.abs(w_single - w_parallel)):.2e}")
   

3. Cài đặt một tầng Hỗn hợp chuyên gia đơn giản. Tạo một mạng cổng định tuyến các token đến top-K chuyên gia và kết hợp đầu ra của chúng.

   import jax
   import jax.numpy as jnp
   
   def expert_fn(x, W1, b1, W2, b2):
       """Một chuyên gia FFN 2 lớp đơn giản."""
       h = jnp.maximum(0, x @ W1 + b1)  # ReLU
       return h @ W2 + b2
   
   def moe_layer(x, gate_W, experts_params, top_k=2):
       """
       Lượt truyền xuôi của MoE.
       x: (batch, d_model)
       gate_W: (d_model, n_experts)
       experts_params: danh sách (W1, b1, W2, b2) cho mỗi chuyên gia
       """
       n_experts = len(experts_params)
   
       # Cổng: tính điểm định tuyến
       gate_logits = x @ gate_W  # (batch, n_experts)
       gate_probs = jax.nn.softmax(gate_logits, axis=-1)
   
       # Chọn top-K
       top_k_indices = jnp.argsort(-gate_probs, axis=-1)[:, :top_k]
       top_k_probs = jnp.take_along_axis(gate_probs, top_k_indices, axis=-1)
       # Chuẩn hóa lại
       top_k_probs = top_k_probs / jnp.sum(top_k_probs, axis=-1, keepdims=True)
   
       # Tính đầu ra chuyên gia (đơn giản hóa: chạy mọi chuyên gia, lọc sau)
       expert_outputs = jnp.stack([
           expert_fn(x, *experts_params[i]) for i in range(n_experts)
       ], axis=1)  # (batch, n_experts, d_model)
   
       # Thu thập đầu ra top-K chuyên gia và trọng hóa chúng
       batch_idx = jnp.arange(x.shape[0])[:, None]
       selected_outputs = expert_outputs[batch_idx, top_k_indices]  # (batch, top_k, d_model)
       output = jnp.sum(selected_outputs * top_k_probs[:, :, None], axis=1)
   
       return output, gate_probs
   
   # Thiết lập
   key = jax.random.PRNGKey(42)
   batch, d_model, d_ff, n_experts = 8, 16, 32, 4
   
   # Khởi tạo chuyên gia
   experts_params = []
   for i in range(n_experts):
       k1, k2, key = jax.random.split(key, 3)[0], jax.random.split(key, 3)[1], jax.random.split(key, 3)[2]
       experts_params.append((
           jax.random.normal(k1, (d_model, d_ff)) * 0.1,
           jnp.zeros(d_ff),
           jax.random.normal(k2, (d_ff, d_model)) * 0.1,
           jnp.zeros(d_model),
       ))
   
   key, subkey = jax.random.split(key)
   gate_W = jax.random.normal(subkey, (d_model, n_experts)) * 0.1
   x = jax.random.normal(key, (batch, d_model))
   
   output, gate_probs = moe_layer(x, gate_W, experts_params, top_k=2)
   
   print(f"Input shape: {x.shape}")
   print(f"Output shape: {output.shape}")
   print(f"Gate probabilities (first sample): {gate_probs[0]}")
   print(f"Expert usage (avg across batch):")
   for i in range(n_experts):
       usage = jnp.mean(gate_probs[:, i])
       print(f"  Expert {i}: {usage:.3f}")