边缘计算环境下NanoGPT分布式推理实践

1. 项目概述：边缘计算环境下的NanoGPT分布式推理

在资源受限的边缘设备上部署大型语言模型一直是行业难题。这个项目将斯坦福CS149课程中的NanoGPT模型移植到树莓派5集群，实现了基于TCP通信的分布式注意力计算。主节点负责拆分计算任务并协调结果，从节点执行部分注意力计算，最终合并输出。这种架构特别适合智能家居网关、工业现场设备等边缘场景，可以在不依赖云端的情况下实现本地化AI推理。

2. 主节点改造方案详解

2.1 模型架构调整

主节点需要同时处理本地计算和远程通信，因此对原始model.py的修改主要集中在注意力机制部分。关键改动包括：

1. 注意力头数量调整：将n_head从12减至4，这是考虑到树莓派的计算能力限制。实验表明，4个注意力头在保持模型质量的同时，能更好地适应边缘设备性能。

python复制@dataclass
class GPTConfig:
    n_head: int = 4  # 修改为4个头
    testname: str = "distributed"  # 默认测试模式设为分布式

2. 分布式计算分支：在CausalSelfAttention的forward函数中添加分布式处理逻辑。这里采用半本地半远程的计算策略，既保证计算效率，又实现负载均衡。

2.2 网络通信实现

主节点与从节点的通信采用TCP协议，通过多线程实现计算与通信的并行：

python复制def compute_remote():
    # 建立TCP连接
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.connect(('192.168.1.100', 8888))
    
    # 发送张量维度信息
    B_r, nh_r, T_r, hs_r = q_remote.size()
    sock.sendall(struct.pack('!4i', B_r, nh_r, T_r, hs_r))
    
    # 发送张量数据
    sock.sendall(q_remote.numpy().tobytes())
    sock.sendall(k_remote.numpy()