import torch.distributed as dist

def setup(rank,local_rank, world_size):
    dist.init_process_group(
        backend='nccl',    # 使用NCCL后端（GPU场景）
        init_method='env://',  # 从环境变量读取配置
        rank=rank,
        world_size=world_size
    )
    torch.cuda.set_device(local_rank)  # 绑定当前GPU

## pytorch.distribute.launch
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--rank", type=int) 
parser.add_argument("--local_rank", type=int) 
parser.add_argument("--world_size", type=int)
args = parser.parse_args()
rank = args.rank
local_rank = args.local_rank
world_size = args.world_size

## torchrun
import os 
rank = int(os.environ['RANK'] )
local_rank = int(os.environ['LOCAL_RANK'])
world_size = int(os.environ['WORLD_SIZE'])

def setup(rank,local_rank, world_size):
    dist.init_process_group(
        backend='nccl',    # 使用NCCL后端（GPU场景）
    )
    torch.cuda.set_device(local_rank)  # 绑定当前GPU
setup(rank,world_size)

## 数据集合
## 构造
sampler = torch.utils.data.distributed.DistributedSampler(dataset)
data_loader = DataLoader(dataset, batch_size=batch_size, sampler=sampler)

from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

##  必须init_process_group 之后才可以调用 
model.to(device)
model = DDP(model, device_ids=[local_rank])
model = FSDP(model, device_id=local_rank)

## 要在构造DDP model之后，才能用model初始化optimizer。
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

import functools
from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
my_auto_wrap_policy = functools.partial(
        size_based_auto_wrap_policy, min_num_params=20000
    )

from torch.distributed.fsdp import CPUOffload
cpu_offload = CPUOffload(offload_params=True)

from torch.distributed.fsdp.mixed_precision import MixedPrecision
mixed_precision_policy = MixedPrecision(
    param_dtype=torch.float16,  # 模型参数的精度
    buffer_dtype=torch.float16,  # 模型缓冲区的精度
    reduce_dtype=torch.float16,  # 梯度归约的精度
    backward_dtype=torch.float16,  # 反向传播的精度
    keep_low_precision_grads=True  # 是否保持梯度的低精度
)

from torch.distributed.fsdp import ShardingStrategy
sharding_strategy = ShardingStrategy.FULL_SHARD 
model = FSDP(
    model,
    auto_wrap_policy=auto_wrap_policy,
    sharding_strategy=sharding_strategy,
    device_id=torch.cuda.current_device(),
)

param_init_fn = lambda module: module.to_empty(device=torch.device("cuda"), recurse=False)

if dist.get_rank() == 0:
    model.load_state_dict(torch.load(ckpt_path))
if dist.get_rank() == 0:
	torch.save(model.module.state_dict(), "%d.ckpt" % epoch)

# DDP init
dist.init_process_group(backend='nccl')

# 按照原来的方式定义模型，这里的BN都使用普通BN就行了。
model = MyModel()
# 引入SyncBN，这句代码，会将普通BN替换成SyncBN。
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)

# 构造DDP模型
model = DDP(model, device_ids=[local_rank], output_device=local_rank)

model = DDP(model)

for 每次梯度累加循环
    optimizer.zero_grad()
    # 前accumulation_step-1个step，不进行梯度同步，累积梯度。
    for _ in range(K-1)::
        with model.no_sync():
            prediction = model(data)
            loss = loss_fn(prediction, label) / K
            loss.backward()  # 积累梯度，不应用梯度改变
    # 第K个step，进行梯度同步
    prediction = model(data)
    loss = loss_fn(prediction, label) / K
    loss.backward()  # 积累梯度，不应用梯度改变
    optimizer.step()

code_before()
# 在这一步同步
torch.distributed.barrier()
code_after()

if rank == 0:
    do_A()
    torch.distributed.barrier()
else:
    torch.distributed.barrier()
    do_B()

import logging

# 给主要进程（rank=0）设置低输出等级，给其他进程设置高输出等级。
logging.basicConfig(level=logging.INFO if rank in [-1, 0] else logging.WARN)
# 普通log，只会打印一次。
logging.info("This is an ordinary log.")
# 危险的warning、error，无论在哪个进程，都会被打印出来，从而方便debug。
logging.error("This is a fatal log!")

def main():
    rank = torch.distributed.get_rank()
    # 问题完美解决！
    init_seeds(1 + rank)

for epoch in iterator:
    # DDP：设置sampler的epoch，
    # DistributedSampler需要这个来指定shuffle方式，
    # 通过维持各个进程之间的相同随机数种子使不同进程能获得同样的shuffle效果。
    data_loader.sampler.set_epoch(epoch)

model_engine, optimizer, _, scheduler = deepspeed.initialize(
    config="path",

deepspeed_config = {
    "train_batch_size": 8,
    "gradient_accumulation_steps": 4,  # 
    "steps_per_print": 2,            # 新增关键参数
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": 0.001
        }
    },
    "pipeline": {
        "activation_checkpoint_interval": 1  # 启用激活检查点
    },
    "fp16": {  # 可选：添加混合精度支持
        "enabled": True
    }
}
model_engine, optimizer, _, scheduler = deepspeed.initialize(
    config=deepspeed_config,

    model.to(device)
    parameters = filter(lambda x:x.requires_grad,model.parameters())
    model_engine, optimizer, _, scheduler = deepspeed.initialize(
    config=cfg.train.deepspeed_config,
    model=model,
    model_parameters=parameters
)

## 数据集合
## 构造
sampler = torch.utils.data.distributed.DistributedSampler(dataset)
data_loader = DataLoader(dataset, batch_size=batch_size, sampler=sampler)

client_sd = {"step": step, "epoch": epoch}
model_engine.save_checkpoint(args.save_dir, ckpt_id=step, client_sd=client_sd)

# Load checkpoint
_, client_sd = model_engine.load_checkpoint(args.load_dir, args.ckpt_id)

{
    "train_batch_size": 32,
    "gradient_accumulation_steps": 2,
    "fp16": {
        "enabled": true
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": 0.001,
            "betas": [0.9, 0.999],
            "eps": 1e-08,
            "weight_decay": 0.01
        }
    },
    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": 0.0,
            "warmup_max_lr": 0.001,
            "warmup_num_steps": 100
        }
    },
    "zero_optimization": {
        "stage": 2,
        "contiguous_gradients": true,
        "reduce_scatter": true,
        "allgather_partitions": true
    }
}

"fp16": {
  "enabled": true,
  "loss_scale": 0,
  "loss_scale_window": 1000,
  "hysteresis": 2,
  "min_loss_scale": 1
}
"bf16": { "enabled": true }

import torch
import deepspeed

# Define a simple neural network model
class SimpleModel(torch.nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc1 = torch.nn.Linear(784, 128)
        self.fc2 = torch.nn.Linear(128, 10)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

# Initialize DeepSpeed configuration
deepspeed_config = {
    "train_batch_size": 64,
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": 0.001
        }
    }
}

# Initialize model
model = SimpleModel()

# Initialize DeepSpeed for distributed data parallelity
model_engine, optimizer, _, _ = deepspeed.initialize(
    config=deepspeed_config,
    model=model
)

# Dummy data
inputs = torch.randn(64, 784)
labels = torch.randint(0, 10, (64,))

# Forward pass
outputs = model_engine(inputs)
loss = torch.nn.functional.cross_entropy(outputs, labels)

# Backward pass and optimization
model_engine.backward(loss)
model_engine.step()

import torch
import deepspeed
from deepspeed.pipe import PipelineModule, LayerSpec
import os

# 注意：华为昇腾（Ascend）芯片需使用'hccl'后端，NVIDIA GPU使用'nccl'
deepspeed.init_distributed(dist_backend='hccl')  # 假设使用NVIDIA GPU

# DeepSpeed配置需添加流水线并行参数
deepspeed_config = {
    "train_batch_size": 8,
    "gradient_accumulation_steps": 4,  # 
    "steps_per_print": 2,            # 新增关键参数
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": 0.001
        }
    },
    "pipeline": {
        "activation_checkpoint_interval": 1  # 启用激活检查点
    },
    "fp16": {  # 可选：添加混合精度支持
        "enabled": True
    }
}
# 修改数据集以包含标签
class SimpleDataset(torch.utils.data.Dataset):
    def __init__(self, data_size=1000, input_dim=784, output_dim=10):
        self.data = torch.randn(data_size, input_dim)
        self.labels = torch.randn(data_size, output_dim)
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return (self.data[idx], self.labels[idx])  # 返回输入和标签的元组

# 定义模型层
class SimpleLayer(torch.nn.Module):
    def __init__(self, input_size, output_size):
        super().__init__()
        self.fc = torch.nn.Linear(input_size, output_size)
    
    def forward(self, x):
        return torch.relu(self.fc(x))

# 构建流水线模型
layers = [
    LayerSpec(SimpleLayer, 784, 128),
    LayerSpec(SimpleLayer, 128, 10)
]

pipeline_model = PipelineModule(
    layers=layers,
    loss_fn=torch.nn.CrossEntropyLoss(),
    num_stages=2,          # 流水线阶段数需等于GPU数
    partition_method='uniform',  # 均匀划分层到各个阶段
)
dataset = SimpleDataset(data_size=10240)

# 初始化DeepSpeed引擎
model_engine, optimizer, _, _ = deepspeed.initialize(
    config=deepspeed_config,
    model=pipeline_model,
    model_parameters=pipeline_model.parameters(),
    training_data=dataset
)

# 准备数据加载器

for step in range(100): 
    loss = model_engine.train_batch()

# 初始化DeepSpeed引擎
model_engine, optimizer, _, _ = deepspeed.initialize(
    config=deepspeed_config,
    model=pipeline_model,
    model_parameters=pipeline_model.parameters(),
)

# 准备数据加载器
from deepspeed.utils import RepeatingLoader
datasetloader = torch.utils.data.DataLoader(dataset=dataset,batch_size=deepspeed_config["train_batch_size"])
dataloader = RepeatingLoader(dataloader)  # 转为无限循环的迭代器
data_iter = iter(dataloader)
for step in range(100): 
    loss = model_engine.train_batch(data_iter=dataiter)