Burn 分布式计算实战：集体张量操作、DDP 训练与远程设备执行

Burn 分布式计算实战集体张量操作、DDP 训练与远程设备执行【免费下载链接】burnBurn is a next generation tensor library and Deep Learning Framework that doesnt compromise on flexibility, efficiency and portability.项目地址: https://gitcode.com/GitHub_Trending/bu/burn本篇技术指南围绕 Burn 深度学习框架的分布式计算能力展开涵盖burn::tensor::distributed提供的集体张量操作、基于ExecutionStrategy::ddp的分布式数据并行DDP训练以及通过 Burn 计算服务器WebSocket 与 Iroh 两种传输实现远程设备执行。读完本文你将掌握如何在 Burn 中枚举多设备、同步梯度、将训练透明地迁移到远程 GPU 服务器并理解DistributedContext的生命周期管理与 collective 同步的底层原理。Burn 将分布式能力设计为三个可独立使用、也可自由组合的层次分布式张量 API 提供跨设备组的集体操作原语ExecutionStrategy::ddp借助这些原语在训练循环中同步梯度远程Device则把普通张量运算转发给另一个进程托管的计算服务器远程设备集合同样可以参与 DDP。分布式能力总览能力入口用途集体张量操作burn::tensor::distributed在设备组之间执行 all-reduce 等集体运算DDP 训练burn::train::ExecutionStrategy::ddp在训练循环中同步各设备副本的梯度远程设备Device::remote_websocket/Device::remote_iroh把张量运算转发给另一进程中的 Burn 计算服务器三种能力可以独立使用也可以叠加远程设备集合参与 DDP 即是一种典型组合见下文“DDP on Remote Devices”。分布式张量操作分布式张量 API 目前以 all-reduce 为核心主要由以下类型与函数构成类型或函数作用DistributedContext启动并持有设备组的通信资源DistributedConfig配置设备组的梯度聚合方式ReduceOperation选择Sum或Mean归约all_reduce将张量在所有参与设备上归约并把结果返回给每个设备CollectiveTensor表示尚未同步的集体运算结果常规使用前必须先同步在发起任何 collective 调用之前必须先创建DistributedContext。上下文的 Drop 会关闭其通信服务器因此只要设备组仍处于活跃状态就必须保持上下文存活use burn::tensor::{ Device, DeviceType, Tensor, distributed::{ CollectiveTensor, DistributedConfig, DistributedContext, ReduceOperation, all_reduce, }, }; let devices Device::enumerate(DeviceType::Cuda).into_vec(); let _context DistributedContext::init( devices.clone(), DistributedConfig { all_reduce_op: ReduceOperation::Mean, }, ); // Every participant submits its local tensor with the same device list. let local_tensors: VecTensor2 devices .iter() .map(compute_local_value) .collect(); let collectives: Vec_ local_tensors .into_iter() .map(|tensor| all_reduce(tensor, ReduceOperation::Sum, devices.clone())) .collect(); let reduced: VecTensor2 collectives .into_iter() .map(CollectiveTensor::resolve) .collect();生命周期与实现细节从源码看DistributedContext本质上是多设备同步的资源句柄。其init会调用Dispatch::start_communication_server为传入设备启动底层分布式通信服务器见 crates/burn-tensor/src/tensor/distributed.rspub fn init(devices: VecDevice, config: DistributedConfig) - Self { let dispatch_devices devices .iter() .map(|d| d.as_dispatch().clone()) .collect::Vec_(); Dispatch::start_communication_server(dispatch_devices, config); Self { devices } }而Drop实现则会调用Dispatch::close_communication_server做干净的网络资源拆除impl Drop for DistributedContext { fn drop(mut self) { if !self.devices.is_empty() { Dispatch::close_communication_server(self.devices[0].as_dispatch()); } } }因此示例中的let _context ...并不是随手一写——只要_context未离开作用域通信服务器就会一直存活一旦被丢弃设备组的所有分布式通道随之关闭。同步resolve与assume_resolvedall_reduce返回的是CollectiveTensor而非普通Tensor原因是集体通信可能是异步的。必须调用resolve()来同步该 collective 并取回普通Tensor只有自行安排了同步顺序的代码才应使用unsafe的assume_resolved()pub fn resolve(self) - TensorD { Dispatch::sync_collective(self.handle.device()); Tensor::new(BridgeTensor::float(self.handle)) } pub unsafe fn assume_resolved(self) - TensorD { Tensor::new(BridgeTensor::float(self.handle)) }在all_reduce的实现内部它会从每个设备取 dispatch id调用Dispatch::all_reduce后先用assume_resolved包装成CollectiveTensor把“确保同步”的责任推迟到用户显式调用resolve的时刻见 crates/burn-tensor/src/tensor/distributed.rs。归约配置DistributedConfig目前只有一个字段all_reduce_op取值来自ReduceOperation枚举见 crates/burn-std/src/distributed.rsSum对参与设备上的张量求和Mean求均值即每个副本都拿到平均梯度。pub enum ReduceOperation { Sum, Mean, } pub struct DistributedConfig { pub all_reduce_op: ReduceOperation, }Mean是 DDP 训练中的常见选择因为它能让每个副本基于同一份平均梯度更新模型。此外所有参与者必须按兼容的顺序、以相同的设备组调用 collective否则分布式状态将不一致。实际应用中普通代码通常不会直接调用all_reduce而是使用更上层的 DDP 训练策略。分布式数据并行DDP训练DDP 在每台设备上保留一份模型副本并把训练输入切分到各副本。每个副本在本地完成前向与反向传播随后 Burn 对梯度做 all-reduce最后才执行优化器更新。当all_reduce_op配置为ReduceOperation::Mean时每个副本都收到平均梯度从而保持模型同步use burn::{ tensor::{Device, DeviceType, distributed::{DistributedConfig, ReduceOperation}}, train::{ExecutionStrategy, Learner, SupervisedTraining}, }; // List all available CUDA devices let devices Device::enumerate(DeviceType::Cuda).into_vec(); let strategy ExecutionStrategy::ddp( devices, DistributedConfig { all_reduce_op: ReduceOperation::Mean, }, ); // Init the model on the main device with autodiff let model ModelConfig::new().init(strategy.main_device().clone().autodiff()); // Launch DDP training let training SupervisedTraining::new(artifact_dir, dataloader_train, dataloader_valid) .with_training_strategy(strategy.into()) .num_epochs(config.num_epochs); let result training.launch(Learner::new(model, optimizer, lr_scheduler));这种写法使模型构造完全独立于所选策略是单设备、多设备还是 DDP。Learner 统一负责模型副本管理、数据分发、梯度集体同步以及DistributedContext的生命周期。ExecutionStrategy::ddp的底层行为从训练策略源码可以看到ExecutionStrategy::ddp本身就是DistributedContext的创建点见 crates/burn-train/src/learner/supervised/strategies/base.rspub fn ddp(devices: VecDevice, config: DistributedConfig) - Self { let context DistributedContext::init(devices.clone(), config); Self::DistributedDataParallel { devices, context } }也就是说策略一旦创建通信上下文即被初始化上下文随策略一同被 Learner 持有训练循环期间始终存活训练结束后随策略释放。DDP 与 MultiDevice 的区别DDP 与ExecutionStrategy::MultiDevice有着本质差异DDP 为每个设备维护一份模型副本并通过 collective 同步梯度而多设备策略走的是 Burn 非 DDP 的多设备训练路径例如基于优化器分片MultiDeviceOptim::OptimSharded协调优化过程。在 examples/text-classification/examples/ag-news-train.rs 中可以看到两者的切换方式// 非 DDP 多设备优化器分片 launch(ExecutionStrategy::MultiDevice( devices.into_vec(), burn::train::MultiDeviceOptim::OptimSharded, )) // DDP每个设备一份模型副本梯度 all-reduce launch(ExecutionStrategy::ddp( devices.into_vec(), DistributedConfig { all_reduce_op: ReduceOperation::Mean, }, ))该示例还展示了枚举到的设备集可通过DeviceConfig统一配置数据类型如f16、flex32例如devices.configure(DeviceConfig::default().float_dtype(ElemType::dtype()))。远程设备远程设备实现与本地 CUDA、WGPU 或 CPU 设备相同的Device接口。张量创建与运算仍使用普通 API但实际执行发生在 Burn 服务器暴露的设备上let device Device::remote_websocket(ws://localhost:3000, 0); let tensor Tensor::2::ones([32, 128], device); let output model.to_device(device).forward(tensor);其中第二个参数index用于从服务器暴露的多块设备中选择一块——同一地址、不同 index 指向同一主机上的不同设备见 crates/burn-tensor/src/device.rs 中remote_websocket的文档说明。WebSocket 传输保留兼容WebSocket 远程设备是为既有部署保留的。Device::remote_websocket(address, index)在构造时就会建立连接源码中device.connect()用于初始化连接并获取设备默认设置并需要启用remote-websocket特性见 crates/burn-tensor/src/device.rs。Iroh 传输新集成首选新的原生集成应优先选择 Iroh 传输。Iroh 以对等身份peer identity标识服务器无需固定 WebSocket 地址服务器通过Channel::Iroh和RemoteSecret暴露本地设备客户端经 Iroh endpoint 连接后获得同一个统一的Devicelet endpoint Endpoint::builder(presets::N0).bind().await?; let device Device::remote_iroh(endpoint, server_id, 0); let tensor Tensor::1::from_floats([1.0, 2.0, 3.0], device); let output tensor.square().sum(); // Executed by the remote server.RemoteSecret应当由系统生成随机值并通过受信任的渠道分发其公开身份。Device::remote_iroh_authorized还支持向启用对等授权PeerAuthorizer的服务器发送应用自定义凭据credential: Vecu8见 crates/burn-tensor/src/device.rs。对于浏览器Wasm目标由于无法建立同步连接需要改用异步构造器remote_iroh_async与remote_iroh_authorized_async。从 crates/burn-remote/README.md 可以看出 Iroh 传输的更多设计客户端应用进程拥有 Iroh endpoint 的配置权身份持久化、中继策略、地址查询与集群发现都留在 Burn 之外RemoteNode是进程级对象多个设备应克隆它共享同一个 Iroh endpoint、对等连接池与多路复用的 QUIC 连接计算端通过server::start_async(Device::cuda(0), Channel::Iroh { node })暴露设备也可通过BURN_REMOTE_ALPN把 Burn 处理器注册进应用自己的 Iroh Router张量在不同 Iroh 计算对等端之间迁移时负载不经过客户端中转目标对等端直接与源对等端建立经认证的流并使用绑定到目标身份、限制下载次数的短期随机能力令牌同一计算对等端上的多个设备之间保留进程内快速路径Iroh 对等端与旧式 WebSocket 对等端之间的张量迁移不受支持。远程 DDP远程执行与 DDP 可以自然组合。text-classification示例枚举远程 WebSocket 服务器托管的每一块设备并把它们交给同一个 DDP 策略见 examples/text-classification/examples/ag-news-train.rs 的remote模块/// Address of the burn-remote server to train against. const ADDRESS: str ws://localhost:3000; #[cfg(feature ddp)] pub fn run() { let mut devices Device::enumerate(DeviceType::remote_websocket(ADDRESS)); devices .configure(DeviceConfig::default().float_dtype(ElemType::dtype())) .unwrap(); crate::launch(ExecutionStrategy::ddp( devices.into_vec(), DistributedConfig { all_reduce_op: ReduceOperation::Mean, }, )); }从 Learner 的视角看本地 DDP 与远程 DDP 使用的是同一个VecDevice。远程设备把计算转发给服务器而分布式上下文负责在选定的服务器设备之间协调梯度 collective。若只想用服务器上的一块设备做单设备训练则取枚举结果中的一个即可示例中为devices.into_vec().pop().unwrap()此时不要重复配置设备否则会触发DeviceError::AlreadyInitialized。如何选择方案当计算需要在别处执行、但不需要数据并行同步时使用单块远程设备当训练进程直接可用多块设备时使用本地 DDP当一台 Burn 服务器向客户端暴露多块加速器时使用“远程设备 DDP”的组合。无论采用哪种方案分布式执行都要求参与设备支持所需的 collective 操作在启用 DDP、远程或对应后端特性如cuda、remote、remote-websocket、ddp的前提下这些能力才能被完整激活。三者按需组合即可在不改变模型定义与训练代码结构的前提下把 Burn 训练从单机单卡平滑扩展到多卡乃至跨进程的 GPU 集群。【免费下载链接】burnBurn is a next generation tensor library and Deep Learning Framework that doesnt compromise on flexibility, efficiency and portability.项目地址: https://gitcode.com/GitHub_Trending/bu/burn创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考