init
This commit is contained in:
435
transformers/examples/3D_parallel.py
Normal file
435
transformers/examples/3D_parallel.py
Normal file
@@ -0,0 +1,435 @@
|
||||
# Copyright 2024 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
""":
|
||||
This script is used to test training a model using Tensor Parallelism and Data Parallelism.
|
||||
|
||||
Usage:
|
||||
export CUDA_VISIBLE_DEVICES=0,1,2,3
|
||||
export CUDA_VISIBLE_DEVICES=4,5,6,7
|
||||
export CUDA_VISIBLE_DEVICES=5,6,7
|
||||
TP_SIZE=2 DP_SIZE=2 torchrun --nproc_per_node=4 --rdzv_endpoint=localhost:29503 examples/3D_parallel.py
|
||||
CP_SIZE=2 DP_SIZE=2 torchrun --nproc_per_node=4 examples/3D_parallel.py
|
||||
CP_SIZE=2 TP_SIZE=2 torchrun --nproc_per_node=4 examples/3D_parallel.py
|
||||
DP_SIZE=2 CP_SIZE=2 TP_SIZE=2 torchrun --nproc_per_node=8 examples/3D_parallel.py
|
||||
|
||||
TP_SIZE=1 CP_SIZE=4 torchrun --nproc_per_node=4 examples/3D_parallel.py
|
||||
TP_SIZE=1 DP_SIZE=4 torchrun --nproc_per_node=4 examples/3D_parallel.py
|
||||
TP_SIZE=4 DP_SIZE=1 torchrun --nproc_per_node=4 --rdzv_endpoint=localhost:29503 examples/3D_parallel.py
|
||||
IGNORE_SANITY=1 CP_SIZE=1 TP_SIZE=1 DP_SIZE=1 torchrun --nproc_per_node=1 --rdzv_endpoint=localhost:29504 examples/3D_parallel.py
|
||||
ocalhost:29504 test_train.py
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
from collections.abc import Iterable
|
||||
from contextlib import nullcontext
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.distributed.checkpoint as dcp
|
||||
import torch.optim as optim
|
||||
import wandb
|
||||
from datasets import load_dataset
|
||||
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
|
||||
from torch.distributed.checkpoint.stateful import Stateful
|
||||
from torch.distributed.device_mesh import DeviceMesh
|
||||
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
|
||||
from torch.distributed.fsdp import ShardingStrategy
|
||||
from torch.distributed.tensor import DTensor
|
||||
from torch.distributed.tensor.experimental import context_parallel
|
||||
from torch.nn.attention import SDPBackend, sdpa_kernel
|
||||
from torch.utils.data import DataLoader
|
||||
from torch.utils.data.distributed import DistributedSampler
|
||||
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
|
||||
# torch.use_deterministic_algorithms(True)
|
||||
torch.backends.cudnn.deterministic = True
|
||||
|
||||
# Set up logging
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
||||
datefmt="%m/%d/%Y %H:%M:%S",
|
||||
level=logging.INFO,
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# from torch.distributed.tensor.experimental._attention import set_rotate_method
|
||||
|
||||
# set_rotate_method("alltoall") # CP rotate shards using all-to-all
|
||||
|
||||
|
||||
def main():
|
||||
tp_size = int(os.environ.get("TP_SIZE", "1"))
|
||||
dp_size = int(os.environ.get("DP_SIZE", "1"))
|
||||
cp_size = int(os.environ.get("CP_SIZE", "1")) # Add CP size configuration
|
||||
sdpa_backend = SDPBackend.FLASH_ATTENTION # For CP
|
||||
# sdpa_backend = SDPBackend.MATH # For CP
|
||||
global_batch_size = 8 # Desired global batch size
|
||||
seq_len = 1024 # Sequence length
|
||||
num_train_steps = 10000 # Number of training steps
|
||||
LR = 1e-5
|
||||
model_name = "HuggingFaceTB/SmolLM2-1.7B"
|
||||
# model_name = "unsloth/Llama-3.2-1B"
|
||||
|
||||
CHECKPOINT_DIR = f"checkpoint_tp{tp_size}_dp{dp_size}_cp{cp_size}"
|
||||
|
||||
# Initialize distributed environment
|
||||
if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
|
||||
dist.init_process_group("nccl")
|
||||
rank = dist.get_rank()
|
||||
world_size = dist.get_world_size()
|
||||
local_rank = int(os.environ["LOCAL_RANK"])
|
||||
torch.cuda.set_device(local_rank)
|
||||
|
||||
assert world_size == tp_size * dp_size * cp_size, (
|
||||
f"World size ({world_size}) must equal TP size ({tp_size}) * DP size ({dp_size}) * CP size ({cp_size})"
|
||||
)
|
||||
|
||||
mesh = torch.arange(world_size).reshape(dp_size, tp_size, cp_size)
|
||||
world_mesh = DeviceMesh(device_type="cuda", mesh=mesh, mesh_dim_names=("dp", "tp", "cp"))
|
||||
tp_mesh = world_mesh["tp"]
|
||||
dp_mesh = world_mesh["dp"]
|
||||
cp_mesh = world_mesh["cp"]
|
||||
world_mesh["dp", "cp"]._flatten(mesh_dim_name="dp_cp")
|
||||
logger.info(f"Created DeviceMesh: {world_mesh}")
|
||||
logger.info(
|
||||
f"Distributed setup - Rank: {rank}, World size: {world_size}, Local rank: {local_rank}, DP: {dp_mesh.get_local_rank()}, TP: {tp_mesh.get_local_rank()}, CP: {cp_mesh.get_local_rank()}"
|
||||
)
|
||||
|
||||
if dist.get_rank() == 0:
|
||||
wandb.init(
|
||||
project="tp_dp_test",
|
||||
config={
|
||||
"tp_size": tp_size,
|
||||
"dp_size": dp_size,
|
||||
"cp_size": cp_size,
|
||||
"global_batch_size": global_batch_size,
|
||||
"model_name": model_name,
|
||||
"dataset": "roneneldan/TinyStories-1M",
|
||||
"seq_len": seq_len,
|
||||
"lr": LR,
|
||||
"weight_decay": 0.1,
|
||||
},
|
||||
name=f"llama_tp{tp_size}_dp{dp_size}_cp{cp_size}"
|
||||
if model_name == "unsloth/Llama-3.2-1B"
|
||||
else f"tp{tp_size}_dp{dp_size}_cp{cp_size}",
|
||||
)
|
||||
logger.info("Wandb initialized.")
|
||||
# Log the current file to wandb
|
||||
wandb.save("test_train.py")
|
||||
|
||||
# Load model and tokenizer
|
||||
logger.info(f"Loading model and tokenizer from {model_name}")
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||||
if tokenizer.pad_token is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
logger.info(f"Set pad_token to eos_token: {tokenizer.pad_token}")
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_name,
|
||||
device_mesh=tp_mesh if dist.is_initialized() else None,
|
||||
tp_plan="auto",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
logger.info(f"Model loaded onto device mesh: {tp_mesh}")
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
logger.info(f"Using device: {device} for non-model tensors")
|
||||
use_ddp = False
|
||||
if dist.is_initialized() and dp_mesh.size() > 1:
|
||||
model = FSDP(model, device_mesh=dp_mesh, sharding_strategy=ShardingStrategy.NO_SHARD)
|
||||
use_ddp = True
|
||||
pass
|
||||
|
||||
model.train()
|
||||
|
||||
logger.info("Loading TinyStories dataset...")
|
||||
raw_dataset = load_dataset("roneneldan/TinyStories", split="train[:1%]") # Use 1% for faster testing
|
||||
|
||||
def tokenize_function(examples):
|
||||
# Tokenize the text without padding
|
||||
tokenized_batch = tokenizer(
|
||||
examples["text"], padding=False, truncation=True, max_length=seq_len, return_tensors=None
|
||||
)
|
||||
# Set labels to be the same as input_ids for Causal LM
|
||||
tokenized_batch["labels"] = tokenized_batch["input_ids"].copy()
|
||||
return tokenized_batch
|
||||
|
||||
tokenized_dataset = raw_dataset.map(tokenize_function, batched=True, remove_columns=["text"])
|
||||
logger.info(f"Dataset loaded and tokenized. Size: {len(tokenized_dataset)}")
|
||||
|
||||
# Create packed sequences
|
||||
def create_packed_sequences(examples):
|
||||
# Flatten all sequences
|
||||
all_tokens = []
|
||||
for input_ids in examples["input_ids"]:
|
||||
all_tokens.extend(input_ids)
|
||||
|
||||
# Split into sequences of seq_len + 1 (for input + label)
|
||||
num_sequences = len(all_tokens) // (seq_len + 1)
|
||||
packed_input_ids = []
|
||||
packed_labels = []
|
||||
|
||||
for i in range(num_sequences):
|
||||
start_idx = i * (seq_len + 1)
|
||||
end_idx = start_idx + (seq_len + 1)
|
||||
# Get the full sequence
|
||||
full_sequence = all_tokens[start_idx:end_idx]
|
||||
# For input_ids, remove the last token
|
||||
packed_input_ids.append(full_sequence[:-1])
|
||||
# For labels, remove the first token
|
||||
packed_labels.append(full_sequence[1:])
|
||||
|
||||
return {"input_ids": packed_input_ids, "labels": packed_labels}
|
||||
|
||||
# Apply packing to the dataset
|
||||
packed_dataset = tokenized_dataset.map(
|
||||
create_packed_sequences,
|
||||
batched=True,
|
||||
remove_columns=tokenized_dataset.column_names,
|
||||
batch_size=1000, # Process in batches for efficiency
|
||||
num_proc=60,
|
||||
)
|
||||
logger.info(f"Dataset packed. New size: {len(packed_dataset)}")
|
||||
|
||||
# Shuffle the packed dataset
|
||||
packed_dataset = packed_dataset.shuffle(seed=42)
|
||||
logger.info("Packed dataset shuffled")
|
||||
|
||||
# Calculate local batch size
|
||||
if dist.is_initialized():
|
||||
assert global_batch_size % dp_mesh.size() == 0, (
|
||||
f"Global batch size ({global_batch_size}) must be divisible by DP size ({dp_mesh.size()})"
|
||||
)
|
||||
local_batch_size = global_batch_size // dp_mesh.size()
|
||||
else:
|
||||
local_batch_size = global_batch_size
|
||||
|
||||
logger.info(
|
||||
f"Global batch size: {global_batch_size}, DP size: {dp_size if dist.is_initialized() else 1}, Local batch size: {local_batch_size}"
|
||||
)
|
||||
|
||||
# Simple collate function since sequences are already packed
|
||||
def collate_fn(batch):
|
||||
input_ids = torch.tensor([item["input_ids"] for item in batch], dtype=torch.long)
|
||||
labels = torch.tensor([item["labels"] for item in batch], dtype=torch.long)
|
||||
return {"input_ids": input_ids, "labels": labels}
|
||||
|
||||
if dist.is_initialized():
|
||||
sampler = DistributedSampler(
|
||||
packed_dataset, num_replicas=dp_mesh.size(), rank=dp_mesh.get_local_rank(), shuffle=False
|
||||
)
|
||||
else:
|
||||
sampler = None
|
||||
|
||||
dataloader = DataLoader(
|
||||
packed_dataset,
|
||||
batch_size=local_batch_size,
|
||||
sampler=sampler,
|
||||
shuffle=False,
|
||||
collate_fn=collate_fn,
|
||||
pin_memory=True,
|
||||
)
|
||||
logger.info(f"DataLoader created. Distributed: {dist.is_initialized()}")
|
||||
|
||||
optimizer = optim.AdamW(model.parameters(), lr=LR, weight_decay=0.1)
|
||||
|
||||
# Training loop
|
||||
logger.info(f"Starting training for {num_train_steps} steps...")
|
||||
model.train()
|
||||
step = 0
|
||||
while step < num_train_steps:
|
||||
for batch in dataloader:
|
||||
if step >= num_train_steps:
|
||||
break # Exit loop if max steps reached
|
||||
|
||||
# Move batch to appropriate device
|
||||
batch = {k: v.to(device) for k, v in batch.items()}
|
||||
optimizer.zero_grad()
|
||||
|
||||
# Add position_ids to batch before CP sharding
|
||||
batch_size = batch["input_ids"].shape[0]
|
||||
position_ids = torch.arange(0, seq_len, dtype=torch.long, device=device)
|
||||
position_ids = position_ids.unsqueeze(0).expand(batch_size, -1)
|
||||
batch["position_ids"] = position_ids
|
||||
from torch.distributed.tensor.experimental._attention import _cp_options
|
||||
|
||||
_cp_options.enable_load_balance = False
|
||||
|
||||
with sdpa_kernel(sdpa_backend): # TODO: ideally move this to attention implementation
|
||||
cp_context = (
|
||||
nullcontext()
|
||||
if cp_mesh.size() == 1
|
||||
else context_parallel(
|
||||
cp_mesh,
|
||||
buffers=[
|
||||
batch["input_ids"],
|
||||
batch["labels"],
|
||||
batch["position_ids"],
|
||||
],
|
||||
buffer_seq_dims=[1, 1, 1],
|
||||
)
|
||||
)
|
||||
with cp_context:
|
||||
# Pop labels from batch before model forward pass
|
||||
labels = batch.pop("labels")
|
||||
outputs = model(**batch) # [mbs, seq_len/cp]
|
||||
loss = outputs.loss
|
||||
logits = outputs.logits
|
||||
|
||||
# Compute loss with shifted labels
|
||||
loss = model.loss_function(
|
||||
logits=logits, labels=None, shift_labels=labels, vocab_size=model.config.vocab_size
|
||||
)
|
||||
loss.backward()
|
||||
|
||||
# all reduce grads across dp_cp if applicable
|
||||
all_reduce_grads(model, world_mesh, use_ddp=use_ddp)
|
||||
|
||||
if hasattr(model, "clip_grad_norm_"):
|
||||
gradnorm = model.clip_grad_norm_(max_norm=1.0, norm_type=2.0) # TODO: fix reported gradnorm
|
||||
else:
|
||||
# only works with FSDP's NO_SHARD otherwise we should use FSDP's clip_grad_norm_
|
||||
assert len(list(model.parameters())) > 5, "No parameters found in model. Probably DDP bug.."
|
||||
gradnorm = clip_grad_norm_(model.parameters(), max_norm=1.0, norm_type=2.0, foreach=True)
|
||||
|
||||
optimizer.step()
|
||||
# allreduce loss across cp_dp before logging
|
||||
if dist.is_initialized() and (cp_mesh.size() > 1 or dp_mesh.size() > 1):
|
||||
dist.all_reduce(loss, group=world_mesh["dp_cp"].get_group(), op=dist.ReduceOp.AVG)
|
||||
current_loss = loss.item()
|
||||
|
||||
# Log loss and gradnorm to wandb (only on rank 0 of dp group)
|
||||
if not dist.is_initialized() or dist.get_rank() == 0:
|
||||
logger.info(
|
||||
f"Step: {step} | GBS: {global_batch_size} | DP: {dp_mesh.size()} | TP: {tp_mesh.size()} | CP: {cp_mesh.size()} | Loss: {current_loss} | Gradnorm: {gradnorm} | lr: {LR}"
|
||||
)
|
||||
wandb.log(
|
||||
{
|
||||
"train/loss": current_loss,
|
||||
"train/gradnorm": gradnorm,
|
||||
"step": step,
|
||||
"lr": LR,
|
||||
"GBS": global_batch_size,
|
||||
}
|
||||
)
|
||||
|
||||
step += 1 # Increment step count
|
||||
|
||||
logger.info("Training loop finished.")
|
||||
|
||||
# Save model using DCP (only if distributed)
|
||||
if dist.is_initialized():
|
||||
state_dict = {"app": AppState(model, optimizer)}
|
||||
dcp.save(
|
||||
state_dict=state_dict,
|
||||
checkpoint_id=CHECKPOINT_DIR,
|
||||
)
|
||||
logger.info(f"Saved checkpoint to {CHECKPOINT_DIR}")
|
||||
else:
|
||||
# Fallback to regular save for non-distributed case
|
||||
save_dir = "test_model_nondist"
|
||||
model.save_pretrained(save_dir, safe_serialization=False)
|
||||
tokenizer.save_pretrained(save_dir) # Save tokenizer too
|
||||
logger.info(f"Saved model to {save_dir}")
|
||||
|
||||
dist.destroy_process_group()
|
||||
logger.info("Cleaned up distributed process group")
|
||||
# Finish wandb run on rank 0
|
||||
if dist.get_rank() == 0:
|
||||
wandb.finish()
|
||||
logger.info("Wandb run finished.")
|
||||
|
||||
|
||||
def all_reduce_grads(model, world_mesh, use_ddp):
|
||||
"""All reduce gradients across dp_cp if applicable."""
|
||||
cp_mesh = world_mesh["cp"]
|
||||
if use_ddp:
|
||||
# DDP/FSDP takes care of syncing grads
|
||||
mesh = cp_mesh
|
||||
else:
|
||||
mesh = world_mesh["dp", "cp"]._flatten(mesh_dim_name="dp_cp")
|
||||
if dist.is_initialized() and mesh.size() > 1:
|
||||
for name, param in model.named_parameters():
|
||||
if param.grad is not None:
|
||||
# Workaround for cross-mesh communication limitation with DTensor gradients
|
||||
if isinstance(param.grad, DTensor):
|
||||
local_grad = param.grad.to_local()
|
||||
# Ensure grad requires grad for inplace modification checks (might not be needed)
|
||||
# local_grad = local_grad.detach().requires_grad_(True)
|
||||
torch.distributed.all_reduce(local_grad, op=torch.distributed.ReduceOp.SUM, group=mesh.get_group())
|
||||
local_grad = local_grad / mesh.size()
|
||||
# Assign averaged grad back - need careful handling if DTensor structure is complex
|
||||
# This simple assignment might work if the grad structure matches param structure
|
||||
param.grad = DTensor.from_local(
|
||||
local_grad, device_mesh=param.grad.device_mesh, placements=param.grad.placements
|
||||
)
|
||||
else:
|
||||
# Handle regular tensors if any exist (e.g. buffers not converted to DTensor)
|
||||
torch.distributed.all_reduce(param.grad, op=torch.distributed.ReduceOp.AVG, group=mesh.get_group())
|
||||
|
||||
|
||||
class AppState(Stateful):
|
||||
"""Wrapper for checkpointing the Application State including model and optimizer."""
|
||||
|
||||
def __init__(self, model, optimizer=None):
|
||||
self.model = model
|
||||
self.optimizer = optimizer
|
||||
|
||||
def state_dict(self):
|
||||
model_state_dict, optimizer_state_dict = get_state_dict(self.model, self.optimizer)
|
||||
return {"model": model_state_dict, "optim": optimizer_state_dict}
|
||||
|
||||
def load_state_dict(self, state_dict):
|
||||
set_state_dict(
|
||||
self.model, self.optimizer, model_state_dict=state_dict["model"], optim_state_dict=state_dict["optim"]
|
||||
)
|
||||
|
||||
|
||||
def clip_grad_norm_(
|
||||
parameters: Iterable[torch.Tensor],
|
||||
max_norm: float,
|
||||
norm_type: float = 2.0,
|
||||
error_if_nonfinite: bool = False,
|
||||
foreach: bool | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Clip the gradient norm of an iterable of parameters.
|
||||
"""
|
||||
# Filter out parameters with no gradients
|
||||
parameters = [p for p in parameters if p.grad is not None]
|
||||
assert len(parameters) > 0, "No parameters with gradients found"
|
||||
|
||||
# Calculate total norm
|
||||
if norm_type == float("inf"):
|
||||
total_norm = max(p.grad.detach().abs().max() for p in parameters)
|
||||
else:
|
||||
total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), norm_type) for p in parameters]), norm_type)
|
||||
|
||||
# Convert DTensor to local tensor if needed
|
||||
if isinstance(total_norm, DTensor):
|
||||
total_norm = total_norm.full_tensor()
|
||||
|
||||
# Clip gradients
|
||||
clip_coef = max_norm / (total_norm + 1e-6)
|
||||
if clip_coef < 1:
|
||||
for p in parameters:
|
||||
p.grad.detach().mul_(clip_coef)
|
||||
|
||||
return total_norm
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user