在 Windows11 + AMD ROCm 环境下运行 LlamaFactory 0.9.4 时,需要伪造缺失的分布式模块。以下是完整的补丁代码,可以保存为 start_llamafactory.py 使用。

import sys
import types

def patch_torch_distributed():
    try:
        import torch.distributed as dist
    except ImportError:
        dist = types.ModuleType("torch.distributed")
        sys.modules["torch.distributed"] = dist

    # 基础类
    class Store: pass
    class ProcessGroup: pass
    class FileStore(Store): pass
    class TCPStore(Store): pass
    class Backend:
        @staticmethod
        def register_backend(name, func, *args, **kwargs): return None

    dist.Store = Store
    dist.ProcessGroup = ProcessGroup
    dist.FileStore = FileStore
    dist.TCPStore = TCPStore
    dist.Backend = Backend

    # 常用函数
    dist.is_available = lambda: False
    dist.init_process_group = lambda *args, **kwargs: None
    dist.destroy_process_group = lambda *args, **kwargs: None
    dist.get_rank = lambda: 0
    dist.get_world_size = lambda: 1
    dist.barrier = lambda *args, **kwargs: None
    dist.all_reduce = lambda *args, **kwargs: None
    dist.broadcast = lambda *args, **kwargs: None

    # 伪造 torch._C._distributed_c10d
    if "torch._C._distributed_c10d" not in sys.modules:
        fake_c10d = types.ModuleType("torch._C._distributed_c10d")

        # 缺失的类和选项
        class _DistributedBackendOptions: pass
        class AllgatherOptions: pass
        class AllreduceCoalescedOptions: pass
        class AllreduceOptions: pass
        class AllToAllOptions: pass
        class BarrierOptions: pass
        class BroadcastOptions: pass
        class DebugLevel: pass
        class GatherOptions: pass
        class ReduceOp: pass
        class GroupMember: pass
        class Work: pass

        # 缺失的函数
        def _register_process_group(*args, **kwargs): return None
        def _resolve_process_group(*args, **kwargs): return None
        def _unregister_all_process_groups(*args, **kwargs): return None
        def _unregister_process_group(*args, **kwargs): return None

        # 注入
        fake_c10d._DistributedBackendOptions = _DistributedBackendOptions
        fake_c10d.AllgatherOptions = AllgatherOptions
        fake_c10d.AllreduceCoalescedOptions = AllreduceCoalescedOptions
        fake_c10d.AllreduceOptions = AllreduceOptions
        fake_c10d.AllToAllOptions = AllToAllOptions
        fake_c10d.BarrierOptions = BarrierOptions
        fake_c10d.BroadcastOptions = BroadcastOptions
        fake_c10d.DebugLevel = DebugLevel
        fake_c10d.GatherOptions = GatherOptions
        fake_c10d.ReduceOp = ReduceOp
        fake_c10d.GroupMember = GroupMember
        fake_c10d.Work = Work
        fake_c10d._register_process_group = _register_process_group
        fake_c10d._resolve_process_group = _resolve_process_group
        fake_c10d._unregister_all_process_groups = _unregister_all_process_groups
        fake_c10d._unregister_process_group = _unregister_process_group

        sys.modules["torch._C._distributed_c10d"] = fake_c10d

    # 伪造 fsdp 模块
    if "torch.distributed.fsdp" not in sys.modules:
        fake_fsdp = types.ModuleType("torch.distributed.fsdp")
        class FlatParameter: pass
        class FullyShardedDataParallel: pass
        fake_fsdp.FlatParameter = FlatParameter
        fake_fsdp.FullyShardedDataParallel = FullyShardedDataParallel
        sys.modules["torch.distributed.fsdp"] = fake_fsdp

    # 伪造 shard 模块
    if "torch.distributed._shard.sharded_tensor.api" not in sys.modules:
        fake_api = types.ModuleType("torch.distributed._shard.sharded_tensor.api")
        class ShardedTensor: pass
        fake_api.ShardedTensor = ShardedTensor
        sys.modules["torch.distributed._shard.sharded_tensor.api"] = fake_api

    if "torch.distributed._shard.sharded_tensor.shard" not in sys.modules:
        fake_shard = types.ModuleType("torch.distributed._shard.sharded_tensor.shard")
        class Shard: pass
        fake_shard.Shard = Shard
        sys.modules["torch.distributed._shard.sharded_tensor.shard"] = fake_shard

    # 伪造 tensor 模块
    if "torch.distributed.tensor" not in sys.modules:
        sys.modules["torch.distributed.tensor"] = types.ModuleType("torch.distributed.tensor")

    print("[Patch] torch.distributed 已被替换为完整 stub(单机模式)")

# 打补丁
patch_torch_distributed()

# 启动 LlamaFactory WebUI
if __name__ == "__main__":
    from llamafactory.cli import main
    main()

使用方法

  1. 将上述代码保存为 start_llamafactory.py
  2. 在同目录下创建一个批处理文件 start_llamafactory.bat
    @echo off
    set USE_TORCH_DISTRIBUTED=0
    set USE_DISTRIBUTED=0
    set TORCH_COMPILE_DISABLE=1
    python start_llamafactory.py webui
    pause
    
  3. 双击 .bat 文件即可启动 WebUI。

这样,读者就能直接复制完整补丁代码,快速解决 Windows ROCm 下 LlamaFactory 0.9.4 的分布式报错问题。