探索与构建 LLaMA 3 架构：深入组件、编码和推理技术

Llama 3解读了 Meta 官方发布的 Llama 3 介绍文章，侧重 Llama 3 的性能，在前几代 Llama 系列上的改进，以及系统化的模型安全方法。

本文从另一个工程化的角度，剖析了 Llama 3 模型架构中用到的组件和算法。更难得的是，作者为每个知识点配备了简洁易懂的基于 Pytorch 的代码实现。

Llama 系列模型从 Meta 开源的第一版本开始，基本已经是开源大语言模型架构的事实标准。无论是 Meta 自己发布的 Llama 后续版本，还是其他公司借鉴优化的版本，对模型架构本身没有太大的改进，更多的是训练数据在量和质上的增强，以及模型效率在工程方面的提升。

本文来自 Medium，文章质量普遍较高，但订阅费用高，不过很多免费文章也非常不错，比如此篇。

下文由大模型翻译原文： Exploring and building the LLaMA 3 Architecture而来。

---

by vignesh yaadav | 2024 年 4 月 | Medium

随着新的人工智能（AI）模型 Llama 3 和新版 Meta AI 的发布，Meta 在 AI 竞赛中提升了其游戏水平。这款由 Llama 3 驱动的虚拟助手现在可在所有 Meta 平台上使用。

以下是关于 Meta 最新的大型语言模型（LLM）和 AI 助理需要了解的所有信息。

什么是 Llama 3？

Meta 发布了 Llama 系列的最新开放源代码 AI 模型 Llama 3。Llama 3 提供两种版本：一个拥有 80 亿参数，另一个拥有 700 亿参数。

Meta 声称 Llama 3 在这些参数规模上为大型语言模型设定了新标准。他们改进了预训练和后训练过程，从而降低了错误拒绝率，提高了对齐性，并使模型产生了更多样化的响应。值得注意的是，Llama 3 在推理、代码生成和指令跟随方面具有增强的能力。

LLaMA 架构：

与前代模型的主要区别在于，LLaMA-2 在 2 万亿个令牌上进行预训练，比前一代增加了 650%，而 LLaMA-3 则在 15 万亿个令牌上进行训练。模型的上下文长度在 8B 和 70B 模型上都从 4K 增加到了 8K。此外，对于 8B 和 70B 变体，LLaMA 采用了分组查询注意力（Grouped-Query Attention，GQA），而之前一代的 GQA 只在较大的 34B 和 70B 模型中使用。我觉得最具影响力的部分是新的安全方法，采用了两个奖励模型来确保安全性和助益性。

LLaMA 3 从其前一代模型中继承了架构

LLaMA 3：

模型大小、架构、优化超参数

Meta 的 GitHub 链接：llama3/MODEL_CARD.md at main · meta-llama/llama3 (github.com)

LLaMa 2：

模型大小、架构、优化超参数

LLaMA 1:

LLaMA 架构

LLaMA 3 的架构主要沿用了 LLaMA 2 的设计，其中在 8B 和 70B 模型中都使用了 GQA（分组查询注意力）。RoPE（旋转位置嵌入）用于 Q 和 K，而 V 只在应用 SoftMax 函数之前进行乘法。RMS（均方根误差）用于在自注意力和前馈块之前进行归一化，KV 缓存也与 LLaMA 中的保持一致。

注意：此模型架构专注于推理，不适用于训练，因此不包括解码器块和交叉注意力。此外，KV 缓存不会用于模型的训练阶段。

现在让我们深入理解这些组件的每一个。

RoPE（旋转位置编码）

在深入了解 RoPE 之前，了解绝对位置编码和相对编码之间的区别很重要。

绝对位置编码是固定向量，添加到令牌的嵌入中以表示其在句子中的绝对位置。所以，它一次处理一个令牌。你可以将其想象为地图上的（纬度，经度）对：地球上的每个点都有一个独特的对。

相对位置编码则一次处理两个令牌，并在计算注意力时发挥作用：因为注意力机制捕捉两个词之间的 “相关性”，相对位置编码告诉注意力机制所涉及的两个词之间的距离。因此，给定两个令牌，我们创建一个表示它们之间距离的向量。

旋转位置编码可以被视为绝对位置嵌入和相对位置嵌入之间的中间地带，每个令牌都有一个固定的或绝对的嵌入值，并与它的极坐标形式进行内点乘，这相对于二维平面上向量的旋转。

• 在注意力机制中使用的点积是一种内积，可以看作是点积的一般化。

• 我们能否找到一种仅依赖于注意力机制中使用的两个向量q（查询）和k（键）以及它们所代表的令牌的相对距离的内积呢？

• 我们可以定义一个仅依赖于两个嵌入向量q和k以及它们之间相对距离的函数g，如下所示：

使用欧拉公式，我们可以将其写成矩阵形式。

在二维空间中的旋转矩阵，因此得名 旋转型 位置嵌入。

来自 Wolfram MathWorld: https://mathworld.wolfram.com/RotationMatrix.html

• 旋转位置嵌入仅应用于查询和键，而不应用于值。
• 旋转位置嵌入在向量q和k乘以注意力机制中的W矩阵之后应用，而在原始变压器中，它们在应用之前就已施加。

def precomputed_theta_pos_frequencies(head_dim: int, seq_len: int, device: str, theta: float = 10000.0):  
	# As written in the paper, the dimentions o the embedding must be even  
	assert head_dim % 2 == 0, "The head_dim must be even"  
	# Built the theta parameters  
	# According to the formula theta_i = 10000 ^ (-2(i-1)/dim) for i = [1,2,3,..dim/2]  
	# Shape: (head_dim / 2)  
	theta_numerator = torch.arange(0, head_dim, 2).float()  
	# Shape : (head_dim / 2)  
	theta = 1.0 / (theta ** (theta_numerator / head_dim)).to(device)  
	# Construct the positions (the "m" parameter)  
	# shape: (seq_len)  
	m = torch.arange(seq_len, device=device)  
	# multiply each theta by each position using the outer product  
	# shape : (seq_len) outer_product * (head_dim / 2) -> (seq_len, head_dim / 2)  
	freq = torch.outer(m, theta).float()  
	# we can computer complex numbers in the polar form c = R * exp(i * m * theta), where R = 1 as follow  
	# shape: (seq_len, head_dim/2) -> (seq-len, head_dim/2)  
	freq_complex = torch.polar(torch.ones_like(freq), freq)  
	return freq_complex  
  
def apply_rotary_embeddings(x: torch.Tensor, freq_complex: torch.Tensor, device: str):  
	# We transform the each subsequent pair of tokens into a pair of complex numbers  
	# shape : (B, seq_len, head_dim) -> (B, seq_len, h, head_dim / 2)  
	x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))  
	# shape : (seq_len, head_dim / 2) -> (1, seq_len, 1, head_dim / 2)  
	freq_complex = freq_complex.unsqueeze(0).unsqueeze(2)  
	# shape : (B, seq_len, h, head_dim / 2) * (1, seq_len, 1, head_dim / 2) = (B, seq_len, h, head_dim / 2)  
	x_rotate = x_complex * freq_complex  
	# (B, seq_len, h, head_dim / 2) -> (B, seq_len, h, head_dim/2 ,2)  
	x_out = torch.view_as_real(x_rotate)  
	# (B, seq_len, h, head_dim/2, 2) -> (B, seq_len, h * head_dim / 2 * 2)  
	x_out = x_out.reshape(*x.shape)  
	return x_out.type_as(x).to(device)

KV 缓存

• 在推理的每一步中，我们只对模型输出的最后一个令牌感兴趣，因为我们已经拥有了之前的那些。然而，模型需要访问所有先前的令牌来决定要输出哪个令牌，因为它们构成了其上下文（或 “提示”）。
• 这是一种让模型在推理过程中对已见过的令牌做更少计算的方法。解决方案就是KV 缓存！

def repeat_kv(x: torch.Tensor, n_rep: int)-> torch.Tensor:  
	batch_size, seq_len, n_kv_heads, head_dim = x.shape  
	if n_rep == 1:  
	return x  
	else:  
	return (  
		# (B, seq_len, n_kv_heads, 1, head_dim)  
		x[:, :, :, None, :]  
		.expand(batch_size, seq_len, n_kv_heads, n_rep, head_dim)  
		.reshape(batch_size, seq_len, n_kv_heads * n_rep, head_dim)  
	)

分组多查询注意力

分组查询注意力（GQA）是多查询注意力和多头注意力的插值。它在保持与多查询注意力相当的速度的同时，实现了与多头注意力相似的质量。自回归解码的标准做法是缓存序列中先前令牌的键和值，以加速注意力计算。

然而，随着上下文窗口或批次大小的增加，多头注意力（MHA）模型中的键值缓存（kv 缓存）的大小相关的内存成本显著增加。多查询注意力（MQA）是一种机制，它只使用单个键值头处理多个查询，从而节省内存并极大地加快了解码器推理的速度。Llama 采用了（GQA）来解决 Transformer 模型自回归解码过程中的内存带宽挑战。主要问题在于 GPU 进行计算的速度比将其移动到内存中的速度快，每次阶段都需要加载解码器权重和注意力键，这消耗了过多的内存。

class SelfAttention(nn.Module):  
	def __init__(self, args: ModelArgs):  
		super().__init__()  
		self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads  
		# Indicates the number of heads for the queries  
		self.n_heads_q = args.n_heads  
		# Indiates how many times the heads of keys and value should be repeated to match the head of the Query  
		self.n_rep = self.n_heads_q // self.n_kv_heads  
		# Indicates the dimentiona of each head  
		self.head_dim = args.dim // args.n_heads  
		  
		self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)  
		self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)  
		self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)  
		self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)  
		  
		self.cache_k = torch.zeros((args.max_batch_size, args.max_seq_len, self.n_kv_heads, self.head_dim))  
		self.cache_v = torch.zeros((args.max_batch_size, args.max_seq_len, self.n_kv_heads, self.head_dim))  
	  
	def forward(self, x: torch.Tensor, start_pos: int, freq_complex: torch.Tensor):  
		batch_size, seq_len, _ = x.shape #(B, 1, dim)  
		# Apply the wq, wk, wv matrices to query, key and value  
		# (B, 1, dim) -> (B, 1, H_q * head_dim)  
		xq = self.wq(x)  
		# (B, 1, dim) -> (B, 1, H_kv * head_dim)  
		xk = self.wk(x)  
		xv = self.wv(x)  
		  
		# (B, 1, H_q * head_dim) -> (B, 1, H_q, head_dim)  
		xq = xq.view(batch_size, seq_len, self.n_heads_q, self.head_dim)  
		xk = xk.view(batch_size, seq_len, self.n_kv_heads, self.head_dim)  
		# (B, 1, H_kv * head_dim) -> (B, 1, H_kv, head_dim)  
		xv = xv.view(batch_size, seq_len, self.n_kv_heads, self.head_dim)  
		  
		# Apply the rotary embeddings to the keys and values  
		# Does not chnage the shape of the tensor  
		# (B, 1, H_kv, head_dim) -> (B, 1, H_kv, head_dim)  
		xq = apply_rotary_embeddings(xq, freq_complex, device=x.device)  
		xk = apply_rotary_embeddings(xk, freq_complex, device=x.device)  
		  
		# Replace the enty in the cache for this token  
		self.cache_k[:batch_size, start_pos:start_pos + seq_len] = xk  
		self.cache_v[:batch_size, start_pos:start_pos + seq_len] = xv  
		  
		# Retrive all the cached keys and values so far  
		# (B, seq_len_kv, H_kv, head_dim)  
		keys = self.cache_k[:batch_size, 0:start_pos + seq_len]  
		values = self.cache_v[:batch_size, 0:start_pos+seq_len]  
		  
		# Repeat the heads of the K and V to reach the number of heads of the queries  
		keys = repeat_kv(keys, self.n_rep)  
		values = repeat_kv(values, self.n_rep)  
		  
		# (B, 1, h_q, head_dim) --> (b, h_q, 1, head_dim)  
		xq = xq.transpose(1, 2)  
		keys = keys.transpose(1, 2)  
		values = values.transpose(1, 2)  
		  
		# (B, h_q, 1, head_dim) @ (B, h_kv, seq_len-kv, head_dim) -> (B, h_q, 1, seq_len-kv)  
		scores = torch.matmul(xq, keys.transpose(2,3)) / math.sqrt(self.head_dim)  
		scores = F.softmax(scores.float(), dim=-1).type_as(xq)  
		  
		# (B, h_q, 1, seq_len) @ (B, h_q, seq_len-kv, head_dim) --> (b, h-q, q, head_dim)  
		output = torch.matmul(scores, values)  
		  
		# (B, h_q, 1, head_dim) -> (B, 1, h_q, head_dim) -> ()  
		output = (output.transpose(1,2).contiguous().view(batch_size, seq_len, -1))  
		return self.wo(output) # (B, 1, dim) -> (B, 1, dim)

RMS（均方根归一化）

RMSNorm 是一种相对较新的归一化技术，由 Biao Zhang 和 Rico Sennrich 在 2019 年提出。与 BN 和 LN 不同，RMSNorm 是基于激活值的均方根来归一化激活值，而不是使用小批量或层的统计信息。这种方法确保了无论小批量大小或特征数量如何，激活值都能保持一致的缩放比例。此外，RMSNorm 还引入了可学习的缩放参数，提供了类似于批归一化的适应性。

注意：与层规范化一样，我们也有一个可学习的参数gamma（在左侧公式中表示为g），它会乘以归一化的值。

class RMSNorm(nn.Module):  
	def __init__(self, dim: int, eps: float = 1e-5):  
		super().__init__()  
		self.eps = eps  
		# The gamma parameter  
		self.weight = nn.Parameter(torch.ones(dim))  
	  
	def _norm(self, x: torch.Tensor):  
		# (B, seq_len, dim) -> (B, seq_len, 1)  
		return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)  
		  
		def forward(self, x: torch.Tensor):  
		# dim : (B, seq_len, dim) -> (B, seq_len, dim)  
		return self.weight * self._norm(x.float()).type_as(x)

SwiGLU 激活函数

SwiGLU 是一种用于深度神经网络的激活函数，它是 GLU（门控线性单元）的一个变体。它通过计算输入的加权和并对其应用非线性函数来计算神经网络中神经元的输出。SwiGLU 的定义涉及使用包含 Swish 函数和张量乘法的数学表达式。SwiGLU 是 GLU 的一个变体，这意味着它的数学概念与 GLU 相同，但 SwiGLU 使用的非线性函数与 GLU 不同。具体来说，SwiGLU 使用了 Swish 函数，这是一种最近提出的激活函数，在某些应用中已被证明优于其他激活函数。

SwiGLU 有几项优点，使其成为神经网络中有用的激活函数。首先，它基于在许多应用中表现良好的 GLU 概念。其次，它使用了 Swish 函数，该函数在某些情况下表现出优于其他激活函数的性能，特别是在与残差连接结合时。第三，由于使用了逐元素乘法，它允许高效的计算。

• 作者比较了在 Transformer 架构的 Feed-Forward 层中使用不同激活函数时 Transformer 模型的性能。

def forward(self, x: torch.Tensor):  
	swish = F.silu(self.w1(x)) # Apply first transformation  
	x_V = self.w3(x)  
	x = swish * x_V # Apply contraction to original dimension  
	x = self.w2(x) # Apply optional additional transformation  
	return x

前馈块

在 Transformer 架构中，前馈网络层扮演着关键角色，通常紧跟在注意力层和归一化之后。前馈网络层由三个线性变换组成。

class FeedForward(nn.Module):  
	def __init__(self, args: ModelArgs):  
		super().__init__()  
		# Assuming 'hidden_dim' is calculated as per your specifications  
		hidden_dim = 4 * args.dim  
		hidden_dim = int(2 * hidden_dim / 3) # Applying your specific transformation  
		if args.ffn_dim_multiplier is not None:  
		hidden_dim = int(args.ffn_dim_multiplier * hidden_dim)  
		#hidden_dim = int(2 * hidden_dim / 3) # Applying your specific transformation  
		hidden_dim = args.multiple_of * ((hidden_dim + args.multiple_of - 1) // args.multiple_of)  
		  
		self.w1 = nn.Linear(args.dim, hidden_dim, bias=False)  
		self.w2 = nn.Linear(hidden_dim, args.dim, bias=False) # This layer seems to be missing in your original setup  
		self.w3 = nn.Linear(args.dim, hidden_dim, bias=False) # Corrected to match checkpoint  
		  
	def forward(self, x: torch.Tensor):  
		swish = F.silu(self.w1(x)) # Apply first transformation  
		x_V = self.w3(x)  
		x = swish * x_V # Apply contraction to original dimension  
		x = self.w2(x) # Apply optional additional transformation  
		return x

在前向传播过程中，输入张量x会经过多层线性变换。第一层变换后应用的SwiGLU激活函数增强了模型的表现力。最终的变换将张量映射回其原始维度。这种独特的 SwiGLU 激活函数与多个 FeedForward 层的组合提升了模型的性能。

编码器块

由于我们只关注模型的推理过程，所以我们只会关注编码器块。

class EncoderBlock(nn.Module):  
	def __init__(self, args: ModelArgs):  
		super().__init__()  
		self.n_heads = args.n_heads  
		self.dim = args.dim  
		self.head_dim = args.dim // args.n_heads  
		  
		self.attention = SelfAttention(args)  
		self.feed_forward = FeedForward(args)  
		  
		# normalize BEFORE the self attention  
		self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)  
		# Normalization BEFORE the feed forward  
		self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)  
		  
		def forward(self, x: torch.Tensor, start_pos: int, freqs_complex: torch.Tensor):  
		# (B, seq_len, dim) + (B, seq_len, dim) -> (B, seq_len, dim)  
		h = x + self.attention.forward(self.attention_norm(x), start_pos, freqs_complex)  
		out = h + self.feed_forward.forward(self.ffn_norm(h))  
		return out

终极变换器模型

现在，让我们将所有这些独立组件堆叠到变换器块中。

class ModelArgs:  
	dim: int = 4096  
	n_layers: int = 32  
	n_heads: int = 32 # Number of heads for the queries  
	n_kv_heads: Optional[int] = None # Number of heads for the keys and values. If None, defaults to n_heads  
	vocab_size: int = -1 # This will be set when we load the tokenizer  
	multiple_of: int = 256  
	ffn_dim_multiplier: Optional[float] = None # If None, defaults to 4.0  
	norm_eps: float = 1e-6 # only the eps value has been modified int llama 3  
	  
	# Needed for KV cache  
	max_batch_size: int = 32  
	max_seq_len: int = 2048  
	  
	device: str = None

class Transformer(nn.Module):  
	def __init__(self, args: ModelArgs) -> None:  
		super().__init__()  
		  
		assert args.vocab_size != -1, "Vocab size must be set"  
		  
		self.args = args  
		self.vocab_size = args.vocab_size  
		self.n_layers = args.n_layers  
		self.tok_embeddings = nn.Embedding(self.vocab_size, args.dim)  
		  
		self.layers = nn.ModuleList()  
		for _ in range(args.n_layers):  
		self.layers.append(EncoderBlock(args))  
		  
		self.norm = RMSNorm(args.dim, eps=args.norm_eps)  
		self.output = nn.Linear(args.dim, self.vocab_size, bias=False)  
		  
		# To precompute the frequencies of the Rotary Positional Encodings  
		self.freqs_complex = precomputed_theta_pos_frequencies(self.args.dim // self.args.n_heads, self.args.max_seq_len * 2, device=self.args.device)  
		  
		def forward(self, tokens: torch.Tensor, start_pos: int):  
		# (B, seq_len)  
		batch_size, seq_len = tokens.shape  
		assert seq_len == 1, "Only one token at a time can be processed"  
		  
		# (B, seq_len) -> (B, seq_len, dim)  
		h = self.tok_embeddings(tokens)  
		  
		# Retrive the pairs (m, theta) corresponding to the positions [start_pos, start_pos + seq_len]  
		freqs_complex = self.freqs_complex[start_pos:start_pos + seq_len]  
		  
		# Consecutively apply all the encoder layers  
		for layer in self.layers:  
		h = layer(h, start_pos, freqs_complex)  
		h = self.norm(h)  
		output = self.output(h).float()  
		return output

推理

为了进行推理，我们需要从 Meta 的 Llama 3 仓库下载模型的权重。

现在，让我们编写推理模型的代码。在对模型进行推理时，有许多可调整的参数需要考虑，包括 top-p、top-k 以及搜索机制（如贪心搜索 / 波束搜索）。为了简单起见，我只实现了贪心搜索。对于波束搜索，你可以参考 GitHub 上的 Llama 3 仓库中的generation.py。

## Inference  
from typing import Optional  
import torch  
import time  
import json  
from pathlib import Path  
from sentencepiece import SentencePieceProcessor  
from tqdm import tqdm  
from model import ModelArgs, Transformer  
  
  
class LLaMA: 
	def __init__(self, model: Transformer, tokenizer: SentencePieceProcessor, model_args: ModelArgs):  
		self.model = model  
		self.tokenizer = tokenizer  
		self.args = model_args  
	  
	@staticmethod  
	def build(checkpoints_dir: str, tokenizer_path: str, load_model: bool, max_seq_len: int, max_batch_size: int, device: str):  
		prev_time = time.time()  
		if load_model:  
		checkpoints = sorted(Path(checkpoints_dir).glob("*.pth"))  
		assert len(checkpoints) > 0, "No checkpoints files found"  
		chk_path = checkpoints[0]  
		print(f'Loaded checkpoint {chk_path}')  
		checkpoint = torch.load(chk_path, map_location="cpu")  
		print(f'Loaded checkpoint in {(time.time() - prev_time):.2f} seconds')  
		prev_time = time.time()  
		  
		with open(Path(checkpoints_dir) / "params.json", "r") as f:  
		params = json.loads(f.read())  
		model_args: ModelArgs = ModelArgs(  
		max_seq_len = max_seq_len,  
		max_batch_size = max_batch_size,  
		device = device,  
		**params  
		)  
		tokenizer = SentencePieceProcessor()  
		tokenizer.load(tokenizer_path)  
		model_args.vocab_size = tokenizer.vocab_size()  
		  
		if device == "cuda":  
		torch.set_default_tensor_type(torch.cuda.HalfTensor)  
		else:  
		torch.set_default_tensor_type(torch.BFloat16Tensor)  
		  
		model = Transformer(model_args).to(device)  
		  
		if load_model:  
		# Remove repe.freq from checkpoint as we are precomputiong the frequencies  
		del checkpoint["rope.freqs"]  
		model.load_state_dict(checkpoint, strict=False)  
		print(f"Loaded state dict in {(time.time() - prev_time):.2f} seconds")  
		  
		return LLaMA(model, tokenizer, model_args)  
	  
	def text_completion(self, prompts: list[str], temperature: float = 0.6, top_p: float = 0.9, max_gen_len: Optional[int] = None):  
		if max_gen_len is None:  
		max_gen_len = self.args.max_seq_len - 1  
		# Convert each prompt into tokens  
		prompt_tokens = [self.tokenizer.encode(prompt, out_type=int, add_bos=True, add_eos=False) for prompt in prompts]  
		# Make sure the batch size is not too large  
		batch_size = len(prompt_tokens)  
		assert batch_size <= self.args.max_batch_size, f"Batch size {batch_size} is too large"  
		max_prompt_len = max(len(prompt) for prompt in prompt_tokens)  
		# Make sure the prompt length is not larger than the maximum seq length  
		assert max_prompt_len < self.args.max_seq_len, f"Prompt length {max_prompt_len} is too large"  
		total_len = min(self.args.max_seq_len, max_gen_len + max_prompt_len)  
		  
		# Create the list that will contain the generated tokens, along with the initial prompt tokens  
		pad_id = self.tokenizer.pad_id()  
		tokens = torch.full((batch_size, total_len), pad_id, dtype=torch.long, device=self.args.device)  
		for k, t in enumerate(prompt_tokens):  
		tokens[k, :len(t)] = torch.tensor(t, dtype=torch.long, device=self.args.device)  
		  
		eos_reached = torch.tensor([False] * batch_size, device=self.args.device)  
		prompt_tokens_mask = tokens != pad_id # True if the token is a prompt token, False otherwise  
		for cur_pos in tqdm(range(1, total_len), desc='Generating tokens'):  
		with torch.no_grad():  
		logits = self.model.forward(tokens[:, cur_pos-1:cur_pos], cur_pos)  
		if temperature > 0:  
		# The temperature is applied BEFORE the softmax  
		probs = torch.softmax(logits[:, -1] / temperature, dim=-1)  
		next_token = self._sample_top_p(probs, top_p)  
		else:  
		# Greedily select the token with the maximum probability  
		next_token = torch.argmax(logits[:, -1], dim=-1)  
		  
		next_token = next_token.reshape(-1)  
		# Only replace the token if it is a padding token  
		next_token = torch.where(prompt_tokens_mask[:, cur_pos], tokens[:, cur_pos], next_token)  
		tokens[:, cur_pos] = next_token  
		# EOS is reached only if we found an EOS token for padding position  
		eos_reached |= (~prompt_tokens_mask[:, cur_pos]) & (next_token == self.tokenizer.eos_id())  
		if all(eos_reached):  
		break  
	  
		out_tokens = []  
		out_text = []  
		for prompt_index, current_prompt_tokens in enumerate(tokens.tolist()):  
			# Cut to the EOS token, if present  
			if self.tokenizer.eos_id() in current_prompt_tokens:  
			eos_idx = current_prompt_tokens.index(self.tokenizer.eos_id())  
			current_prompt_tokens = current_prompt_tokens[:eos_idx]  
			out_tokens.append(current_prompt_tokens)  
			out_text.append(self.tokenizer.decode(current_prompt_tokens))  
		  
		return (out_tokens, out_text)  
	  
	  
		def _sample_top_p(self, probs, p):  
			probs_sort, probs_idx = torch.sort(probs, dim= -1, descending=True)  
			probs_sum = torch.cumsum(probs_sort, dim=-1)  
			mask = probs_sum - probs_sort > p  
			probs_sort[mask] = 0.0  
			probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))  
			next_token = torch.multinomial(probs_sort, num_samples=1)  
			next_token = torch.gather(probs_idx, -1, next_token)  
			return next_token  
	  
	  
if __name__ == '__main__':  
	import os  
	torch.manual_seed(0)  
	prompts = [  
	# Few shot prompt  
	"""Translate English to kananda:  
	water : ನೀರು  
	land : ಭೂಮಿ  
	dusk : ಸಂಜೆ  
	dawn : ಬೆಳಗುವಿಕೆ  
	milk : ಹಾಲು"""  
	# Zero shot prompt  
	"""Tell me if the following person is actually a real person or a fictional character:  
	Name : Vignesh  
	Decision:  
	"""  
	]  
	allow_cuda = True if 'CUDA_VISIBLE_DEVICES' in os.environ else False  
	device = 'cuda' if torch.cuda.is_available() and allow_cuda else 'cpu'  
	  
	model = LLaMA.build(  
		checkpoints_dir='Meta-Llama-3-8B/',  
		tokenizer_path='Meta-Llama-3-8B/tokenizer.model',  
		load_model=True,  
		max_seq_len=1024,  
		max_batch_size=len(prompts),  
		device=device  
	)  
	  
	print('ALL OK')  
	  
	# Inferencing the model  
	print("Inferenceing the model")  
	out_tokens, out_text = model.text_completion(prompts, max_gen_len=64)  
	assert len(out_tokens) == len(prompts)  
	for i in range(len(out_text)):  
		print(f"{out_text[i]}")  
		print('-' * 50)