DeepSeek源码解析（2）

Tensor （张量）的介绍

在计算机科学和机器学习领域，“张量”（Tensor）是一个数学概念，它被用来表示多维数组。在大模型（如深度学习模型）中，张量扮演着核心角色，
具体来说：

• 数据表示：张量用于表示输入数据、模型参数和中间计算结果。例如，在图像处理中，一张图片可以被表示为一个三维张量（高度、宽度、颜色通道数），而在自然语言处理中，一段文本可以被编码为一系列词向量组成的二维张量（句子长度、词向量维度）。

• 线性代数运算：深度学习中的许多操作，如卷积、全连接等，都可以看作是线性代数中的矩阵运算。这些操作在底层通过张量进行计算。

• 并行计算：由于张量的多维特性，它们非常适合并行计算。现代深度学习框架（如TensorFlow和PyTorch）利用GPU或TPU等硬件的并行处理能力来加速张量的运算。

• 梯度传播：在训练深度学习模型时，需要计算损失函数关于模型参数的梯度。这些梯度也是以张量的形式存在的，并用于更新模型参数。

• 灵活性和扩展性：使用张量可以方便地处理不同形状的数据，并允许模型设计者灵活地构建复杂的网络结构。

• 模块化和抽象：深度学习框架提供了对张量的高级抽象，使得开发者无需关心底层的硬件细节和优化细节，可以更专注于模型的设计和实验。

总之，在大模型中，张量是数据的基本单位，它们使得复杂的数学运算变得可行，并支持了深度学习的快速发展。

DeepSeek源代码 fp8_cast_bf16.py

import os
import json
from argparse import ArgumentParser # 用于解析命令行参数。
from glob import glob # 用于文件名模式匹配
from tqdm import tqdm # 这是一个快速、可扩展的Python进度条库。

import torch # 这是一个流行的开源机器学习库。
from safetensors.torch import load_file, save_file # 从safetensors库的torch子模块中导入了两个函数——load_file和save_file。这两个函数分别用于加载和保存张量数据到文件中。Safetensors是一种安全存储张量数据的格式。

from kernel import weight_dequant #，“weight_dequant”可能是指“权重去量化”的意思。在深度学习模型中，“量化”是一种减少模型大小的技术，通过将权重值映射到有限数量的离散值来实现。“去量化”则可能是将这些离散值转换回原始连续值的过程。

def main(fp8_path, bf16_path):
    """
    将FP8的权重转换为BF16，并保存转换后的权重。

    该函数从指定目录中读取FP8权重，并将其转换为BF16，
    并将转换后的权重保存到另一个指定目录。它还更新了
    为索引文件建模以反映更改。

    参数:
    fp8_path (str)：包含FP8权重和模型索引文件的目录的路径。
    bf16_path (str)：存放转换后的BF16权重值的路径。

    提出了:
    KeyError：如果权重缺少所需的scale_inv张量。

    注:
    —该函数假设FP8权重存储在safetensor文件中。
    —缓存加载的safetensor文件，优化内存使用。
    -该函数更新模型索引文件以删除对scale_inv张量的引用。

    """
    torch.set_default_dtype(torch.bfloat16) # 设置了PyTorch的默认数据类型为bfloat16（Brain Floating Point 16-bit）。这是一种16位浮点数格式，用于深度学习模型中以减少内存占用和加速计算。
    os.makedirs(bf16_path, exist_ok=True) # 创建一个名为bf16_path的目录，如果该目录已经存在，则不会引发错误。
    model_index_file = os.path.join(fp8_path, "model.safetensors.index.json") # 获取模型索引文件的路径。
    with open(model_index_file, "r") as f:  #这行代码打开一个文件用于读取。使用上下文管理器（with语句），这样可以确保文件在读取完毕后会被正确关闭
        model_index = json.load(f)
    weight_map = model_index["weight_map"]
    
    # 缓存加载的安全张量文件
    loaded_files = {}
    fp8_weight_names = []

    # 从正确的文件中获取张量的辅助函数
    def get_tensor(tensor_name):
        """
        从缓存的安全张量文件中检索张量，如果没有缓存，则从磁盘加载。
        参数:
        tensor_name (str)：要检索的张量的名称。

        返回:
            torch.Tensor: 检索到的张量.

        提出了:
            KeyError: 如果该张量在safetensor文件中不存在. 

        """
        file_name = weight_map[tensor_name] # 根据tensor_name（张量的名称）作为键来从weight_map字典中获取对应的文件名file_name
        if file_name not in loaded_files: #检查变量file_name是否不在另一个名为loaded_files的字典中。这个字典可能用来存储已经加载过的文件，以避免重复加载相同的文件。
            file_path = os.path.join(fp8_path, file_name) # 拼接文件路径
            loaded_files[file_name] = load_file(file_path, device="cuda") #使用一个名为load_file的函数来加载权重文件
            # load_file函数接受两个参数：一个是上面构造出的完整文件路径file_path，另一个是设备类型（在这里指定为"cuda"），意味着要将数据加载到GPU上。
            # 加载后的权重被存储在字典loaded_files[file_name]中。
            # 返回已经加载到内存中的、对应于输入参数中指定张量名称（即输入参数中的变量名）的权重数据
        return loaded_files[file_name][tensor_name]
        # 总结来说，这段代码的作用是：根据提供的张量名称查找对应的权重文件，并检查该文件是否已经被加载过；如果没有，则从指定路径加载该权重文件，并将其存储在内存中；
        # 最后返回请求的张量的权重数据。这样的设计可以提高程序效率，避免重复读取相同的数据到内存中。
   
    safetensor_files = list(glob(os.path.join(fp8_path, "*.safetensors"))) #用于查找特定目录下所有以.safetensors为后缀的文件，
    # glob：这是一个Python标准库中的模块，用于从目录通配符搜索中生成文件列表
    # list(...)：将上述函数返回的结果（可能是一个迭代器）转换为一个列表，这样你就可以使用列表的方法
    safetensor_files.sort() #进行排序

     """ 遍历一个包含多个安全张量文件的列表，为每个文件生成一个进度条，提取出每个文件的状态字典，
     并将其存储在一个以文件名为键的大字典中。这样的操作通常用于批量处理模型或数据集，并准备它们以供后续使用。
     """
    for safetensor_file in tqdm(safetensor_files):  # tqdm是一个进度条库，用于在控制台显示循环的进度。
        file_name = os.path.basename(safetensor_file) # 使用os.path.basename()函数从完整的文件路径中提取出文件名（不包括目录路径），并将结果赋值给变量file_name。
        current_state_dict = load_file(safetensor_file, device="cuda") # 是加载文件，并返回一个状态字典（state dictionary），通常在深度学习中用于保存模型的参数。这个状态字典被赋值给变量current_state_dict。
        loaded_files[file_name] = current_state_dict
        
        new_state_dict = {} # 创建一个空字典，用于存储新的状态字典。
        for weight_name, weight in current_state_dict.items(): # 遍历当前状态字典中的每个键值对，其中键是权重名称，值是权重张量。
            if weight_name.endswith("_scale_inv"): #构造逆缩放因子的名称，即在原始权重名称后添加"_scale_inv"。
                continue
            elif weight.element_size() == 1:  # FP8 weight  # 如果权重元素的大小为1字节，则认为这是一个FP8量化的权重。
                scale_inv_name = f"{weight_name}_scale_inv"
                try:
                    # Get scale_inv from the correct file
                    scale_inv = get_tensor(scale_inv_name) # 获取逆缩放因子的实际值。
                    fp8_weight_names.append(weight_name) # 将FP8量化的权重名称添加到列表
                    new_state_dict[weight_name] = weight_dequant(weight, scale_inv) # 将量化的权重和逆缩放因子作为参数传入，进行反量化操作，并将结果存储在新的字典中。
                except KeyError:  #如果在尝试获取逆缩放因子时发生KeyError异常（即找不到对应的逆缩放因子），则打印一条警告信息，并直接将原始的量化权重赋值给新的字典项。
                    print(f"Warning: Missing scale_inv tensor for {weight_name}, skipping conversion")
                    new_state_dict[weight_name] = weight
            else: # 不是FP8量化的权重直接使用原值
                new_state_dict[weight_name] = weight
                
        new_safetensor_file = os.path.join(bf16_path, file_name) # 将多个路径组件合并成一个完整的路径。
        save_file(new_state_dict, new_safetensor_file) # 数据保存到指定的文件路径
        
        # Memory management: keep only the 2 most recently used files 内存管理：只保留最近使用的2个文件
        if len(loaded_files) > 2:  # 检查loaded_files这个列表的长度是否大于2。如果大于2
            oldest_file = next(iter(loaded_files)) # 取出列表中的第一个元素，并将其赋值给变量oldest_file
            # 使用next()函数来获取迭代器的下一个元素，也就是列表的第一个元素。
            del loaded_files[oldest_file] # 删除loaded_files字典中键为oldest_file的项
            torch.cuda.empty_cache() # 清空CUDA缓存。在深度学习训练过程中，GPU内存可能会被占用并积累下来，导致内存泄漏。调用这个函数可以释放不再使用的GPU内存空间
    
    # Update model index
    """
    在一个给定目录下创建或覆盖一个名为“model.safetensors.index.json”的文件，
    并写入一个包含修改后的权重映射信息和空元数据的对象。在这个过程中，它会从权重映射中移除所有以特定模式命名的条目
    （即那些带有后缀“_scale_inv”的条目）
    """
    new_model_index_file = os.path.join(bf16_path, "model.safetensors.index.json") # 拼接路径，生成一个新的文件路径
    for weight_name in fp8_weight_names:
        scale_inv_name = f"{weight_name}_scale_inv"
        if scale_inv_name in weight_map:
            weight_map.pop(scale_inv_name)
    with open(new_model_index_file, "w") as f:
        json.dump({"metadata": {}, "weight_map": weight_map}, f, indent=2)
        

if __name__ == "__main__":  # 检查当前脚本是否作为主程序运行
    parser = ArgumentParser() # ArgumentParser对象用于处理命令行参数。
    parser.add_argument("--input-fp8-hf-path", type=str, required=True) # 参数
    parser.add_argument("--output-bf16-hf-path", type=str, required=True) # 参数
    args = parser.parse_args() #  解析命令行输入的参数，并返回一个包含这些参数值的对象
    main(args.input_fp8_hf_path, args.output_bf16_hf_path) # 调用了名为main的函数，
    

这个文件代码主要功能是将FP8（Float 8-bit）量化的权重转换为BF16（Bfloat 16-bit）格式，并更新模型索引文件以反映这一变化。以下是代码的主要步骤和作用：

1. 设置默认数据类型：将PyTorch的默认数据类型设置为bfloat16。

2. 创建输出目录：如果输出目录不存在，则创建它。

3. 读取模型索引文件：从FP8权重所在的目录中读取模型索引文件，该文件包含了权重映射信息。

4. 缓存加载的安全张量文件：为了避免重复加载相同的文件，代码使用一个字典来缓存已经加载的安全张量文件。

5. 遍历安全张量文件：遍历FP8权重目录下的所有.safetensors文件，并为每个文件生成一个进度条。

6. 提取和转换权重：对于每个安全张量文件，提取出状态字典，并检查每个权重是否是FP8量化的。如果是，使用weight_dequant函数进行去量化操作，并将结果存储在新的状态字典中。

7. 保存转换后的权重：将新的状态字典保存到BF16格式的文件中，并存储在指定的输出目录。

8. 内存管理：为了优化内存使用，代码只保留最近使用的两个安全张量文件，并在不再需要时清空CUDA缓存。

9. 更新模型索引：最后，代码更新模型索引文件，移除所有以_scale_inv结尾的条目（这些是FP8量化时使用的逆缩放因子），并将更新后的索引文件保存到输出目录。

10. 命令行参数解析：脚本接受两个命令行参数：输入FP8权重的路径和输出BF16权重的路径。通过ArgumentParser解析这些参数，并传递给main函数执行上述操作。

总结来说，这个脚本的作用是将FP8量化的深度学习模型权重转换为BF16格式，并更新相关的模型索引文件，以便可以在支持BF16的数据类型上使用这些权重。这有助于减少模型大小并可能提高计算效率。