Unlocking LLaMA2: Key Architecture Insights and Deployment Tricks

Course Recap

In the recent MindSpore public class we explored the principles of LLaMA2 and its inference deployment, demonstrating code examples and discussing state‑of‑the‑art (SOTA) models.

1.1 LLaMA2 vs. Transformer Structure

1.2 Changes from LLaMA1 to LLaMA2

Increased training data

Extended context length

Large‑scale GQA usage

1.3 Detailed LLaMA2 Architecture

RMSNorm : reduces computation of LayerNorm

RoPE : implements complex‑space rotation embeddings

KV‑Cache : addresses performance and memory consumption during inference

Grouped Multi‑Query Attention : balances inference speed and quality

SwiGLU : smoother activation combined with a linear layer

Code example for rotary embeddings:

def apply_rotary_emb(
    xq: mindspore.Tensor,
    xk: mindspore.Tensor,
    freqs_cis: mindspore.Tensor,
) -> Tuple[mindspore.Tensor, mindspore.Tensor]:
    """
    Apply rotary embeddings to input tensors using the given frequency tensor.

    Args:
        xq (mindspore.Tensor): Query tensor.
        xk (mindspore.Tensor): Key tensor.
        freqs_cis (mindspore.Tensor): Precomputed frequency tensor.

    Returns:
        Tuple[mindspore.Tensor, mindspore.Tensor]: Modified query and key tensors.
    """
    xq_ = view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
    xq_out = ops.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = ops.view_as_real(xk_ * freqs_cis).flatten(3)
    return xq_out.astype(xq.dtype), xk_out.astype(xk.dtype)

1.4 Industry LLM Improvement Methods

Performance optimization using FlashAttention, PagedAttention, MQA, GQA

Increasing training data – the most effective way to boost LLM performance

Extending context length to enhance long‑text capabilities

Stabilizing training (e.g., NormHead, Max‑z loss) to prevent loss spikes

Next Session Preview

The upcoming MindSpore public class on January 6 will feature a talk by algorithm engineer Tao Hengtang from Pengcheng Laboratory on the training process of the Pengcheng Brain 200B model, a 2‑trillion‑parameter autoregressive language model trained on the Pengcheng Cloud Brain II platform using MindSpore's multi‑dimensional distributed parallel technology.

Join us at 14:00 for the seventh lecture.