How to Train LLM? - From Data Parallel To Fully Sharded Data Parallel

# Creates model and optimizer in default precision model = Net().cuda() optimizer = optim.SGD(model.parameters(), ...) # Creates a GradScaler once at the beginning of training. scaler = GradScaler() for epoch in epochs: for input, target in data: optimizer.zero_grad() # Runs the forward pass with autocasting. with autocast(device_type='cuda', dtype=torch.float16): output = model(input) loss = loss_fn(output, target) # Scales loss. Calls backward() on scaled loss to create scaled gradients. # Backward passes under autocast are not recommended. # Backward ops run in the same dtype autocast chose for corresponding forward ops. scaler.scale(loss).backward() # scaler.step() first unscales the gradients of the optimizer's assigned params. # If these gradients do not contain infs or NaNs, optimizer.step() is then called, # otherwise, optimizer.step() is skipped. scaler.step(optimizer) # Updates the scale for next iteration. scaler.update()

import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, TensorDataset # Simple network definition class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.fc1 = nn.Linear(10, 5) self.fc2 = nn.Linear(5, 1) def forward(self, x): x = torch.relu(self.fc1(x)) print(f"Activation after fc1 (ReLU) is of type: {x.dtype}") x = self.fc2(x) print(f"Activation after fc2 is of type: {x.dtype}") return x # Create random dataset data = torch.randn(100, 10) # 100 samples, 10 features each targets = torch.randn(100, 1) # 100 target values dataset = TensorDataset(data, targets) dataloader = DataLoader(dataset, batch_size=10, shuffle=True) # Instantiate model and optimizer model = Net().cuda() optimizer = optim.SGD(model.parameters(), lr=0.01) loss_fn = nn.MSELoss() # Training loop for input, target in dataloader: input, target = input.cuda(), target.cuda() # Move data to GPU optimizer.zero_grad() # Enables autocasting for the forward pass (model + loss) with torch.autocast(device_type="cuda"): output = model(input) loss = loss_fn(output, target) print(f"Loss type: {param.dtype}") # Print model parameter precision for name, param in model.named_parameters(): print(f"Parameter {name} is of type: {param.dtype}") # Exits the context manager before backward() loss.backward() # Print gradient precision for name, param in model.named_parameters(): if param.grad is not None: print(f"Gradient for {name} is of type: {param.grad.dtype}") optimizer.step()

Activation after fc1 (ReLU) is of type: torch.float16 Activation after fc2 is of type: torch.float16 Loss type: torch.float32 (Loss Function 거치면서 float32) Parameter fc1.weight is of type: torch.float32 Parameter fc1.bias is of type: torch.float32 Parameter fc2.weight is of type: torch.float32 Parameter fc2.bias is of type: torch.float32 Gradient for fc1.weight is of type: torch.float32 Gradient for fc1.bias is of type: torch.float32 Gradient for fc2.weight is of type: torch.float32 Gradient for fc2.bias is of type: torch.float32

import torch.nn as nn def data_parallel(module, inputs, labels, device_ids, output_device): # (1) [Scatter] data를 device들에 scatter inputs = nn.parallel.scatter(inputs, device_ids) # (2) [Replicate] model weight을 device_ids들에 복제 replicas = nn.parallel.replicate(module, device_ids) # (3) [Forward] 각 device 에 복제된 model이 각 device의 data를 forward logit = nn.parallel.parallel_apply(replicas, inputs) # (4) [Gather] 모델의 logit을 output_device(하나의 device) 로 모음 logit = nn.parallel.gather(outputs, output_device) return logits

import torch import torch.distributed as dist dist.init_process_group("nccl") rank = dist.get_rank() torch.cuda.set_device(rank) tensor = torch.ones(2,2).to(torch.cuda.current_device()) * rank # rank == 0 => [[0, 0], [0, 0]] # rank == 1 => [[1, 1], [1, 1]] # rank == 2 => [[2, 2], [2, 2]] # rank == 3 => [[3, 3], [3, 3]] dist.all_reduce(tensor, op=torch.distributed.ReduceOp.SUM) print(f"rank {rank}: {tensor}\n) # result rank 1: tensor([[6., 6.], [6., 6.]], device='cuda:1') rank 2: tensor([[6., 6.], [6., 6.]], device='cuda:2') rank 0: tensor([[6., 6.], [6., 6.]], device='cuda:0') rank 3: tensor([[6., 6.], [6., 6.]], device='cuda:3')

import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel from torch.optim import AdamW from torch.utils.data import DataLoader, DistributedSampler from transformers import BertForSequenceClassification, BertTokenizer from datasets import load_dataset # 1. initialize process group dist.init_process_group("nccl") rank = dist.get_rank() torch.cuda.set_device(rank) device = torch.cuda.current_device() world_size = dist.get_world_size() # 2. create dataset datasets = load_dataset("multi_nli").data["train"] datasets = [ { "premise": str(p), "hypothesis": str(h), "labels": l.as_py(), } for p, h, l in zip(datasets[2], datasets[5], datasets[9]) ] # 3. create DistributedSampler # DistributedSampler 는 데이터를 쪼개서 다른 프로세스로 전송하기 위한 모듈입니다. sampler = DistributedSampler( datasets, num_replicas=world_size, rank=rank, shuffle=True, ) data_loader = DataLoader( datasets, batch_size=32, num_workers=4, sampler=sampler, shuffle=False, pin_memory=True, ) # 4. create model and tokenizer model_name = "bert-base-cased" tokenizer = BertTokenizer.from_pretrained(model_name) model = BertForSequenceClassification.from_pretrained(model_name, num_labels=3).cuda() # 5. make distributed data parallel module model = DistributedDataParallel(model, device_ids=[device], output_device=device) # 6. create optimizer optimizer = AdamW(model.parameters(), lr=3e-5) # 7. start training for i, data in enumerate(data_loader): optimizer.zero_grad() tokens = tokenizer( data["premise"], data["hypothesis"], padding=True, truncation=True, max_length=512, return_tensors="pt", ) loss = model( input_ids=tokens.input_ids.cuda(), attention_mask=tokens.attention_mask.cuda(), labels=data["labels"], ).loss loss.backward() optimizer.step() if i % 10 == 0 and rank == 0: print(f"step: {i}, loss: {loss}") if i == 300: break

# Copyright (c) Meta Platforms, Inc. and affiliates. # This software may be used and distributed according to the terms of the Llama 2 Community License Agreement. import os import dataclasses import fire import random import torch import torch.optim as optim from peft import get_peft_model, prepare_model_for_kbit_training from torch.distributed.fsdp import ( FullyShardedDataParallel as FSDP, ShardingStrategy ) from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload from torch.optim.lr_scheduler import StepLR from transformers import ( AutoTokenizer, LlamaForCausalLM, LlamaConfig, ) from transformers.models.llama.modeling_llama import LlamaDecoderLayer from llama_recipes.configs import fsdp_config as FSDP_CONFIG from llama_recipes.configs import train_config as TRAIN_CONFIG from llama_recipes.data.concatenator import ConcatDataset from llama_recipes.policies import AnyPrecisionAdamW, apply_fsdp_checkpointing from llama_recipes.utils import fsdp_auto_wrap_policy from llama_recipes.utils.config_utils import ( update_config, generate_peft_config, generate_dataset_config, get_dataloader_kwargs, ) from llama_recipes.utils.dataset_utils import get_preprocessed_dataset from llama_recipes.utils.fsdp_utils import hsdp_device_mesh from llama_recipes.utils.train_utils import ( train, freeze_transformer_layers, setup, setup_environ_flags, clear_gpu_cache, print_model_size, get_policies, ) from accelerate.utils import is_xpu_available def setup_wandb(train_config, fsdp_config, **kwargs): try: import wandb except ImportError: raise ImportError( "You are trying to use wandb which is not currently installed. " "Please install it using pip install wandb" ) from llama_recipes.configs import wandb_config as WANDB_CONFIG wandb_config = WANDB_CONFIG() update_config(wandb_config, **kwargs) init_dict = dataclasses.asdict(wandb_config) run = wandb.init(**init_dict) run.config.update(train_config) run.config.update(fsdp_config, allow_val_change=True) return run def main(**kwargs): # Update the configuration for the training and sharding process train_config, fsdp_config = TRAIN_CONFIG(), FSDP_CONFIG() update_config((train_config, fsdp_config), **kwargs) # Set the seeds for reproducibility if is_xpu_available(): torch.xpu.manual_seed(train_config.seed) torch.manual_seed(train_config.seed) random.seed(train_config.seed) if train_config.enable_fsdp: setup() # torchrun specific local_rank = int(os.environ["LOCAL_RANK"]) rank = int(os.environ["RANK"]) world_size = int(os.environ["WORLD_SIZE"]) if torch.distributed.is_initialized(): if is_xpu_available(): torch.xpu.set_device(local_rank) elif torch.cuda.is_available(): torch.cuda.set_device(local_rank) clear_gpu_cache(local_rank) setup_environ_flags(rank) wandb_run = None if train_config.use_wandb: if not train_config.enable_fsdp or rank==0: wandb_run = setup_wandb(train_config, fsdp_config, **kwargs) # Load the pre-trained model and setup its configuration use_cache = False if train_config.enable_fsdp else None if train_config.enable_fsdp and train_config.low_cpu_fsdp: """ for FSDP, we can save cpu memory by loading pretrained model on rank0 only. this avoids cpu oom when loading large models like llama 70B, in which case model alone would consume 2+TB cpu mem (70 * 4 * 8). This will add some comms overhead and currently requires latest nightly. """ if rank == 0: model = LlamaForCausalLM.from_pretrained( train_config.model_name, load_in_8bit=True if train_config.quantization else None, device_map="auto" if train_config.quantization else None, use_cache=use_cache, attn_implementation="sdpa" if train_config.use_fast_kernels else None, ) else: llama_config = LlamaConfig.from_pretrained(train_config.model_name) llama_config.use_cache = use_cache with torch.device("meta"): model = LlamaForCausalLM(llama_config) else: model = LlamaForCausalLM.from_pretrained( train_config.model_name, load_in_8bit=True if train_config.quantization else None, device_map="auto" if train_config.quantization else None, use_cache=use_cache, attn_implementation="sdpa" if train_config.use_fast_kernels else None, ) # Load the tokenizer and add special tokens tokenizer = AutoTokenizer.from_pretrained(train_config.model_name if train_config.tokenizer_name is None else train_config.tokenizer_name) tokenizer.pad_token_id = tokenizer.eos_token_id # If there is a mismatch between tokenizer vocab size and embedding matrix, # throw a warning and then expand the embedding matrix if len(tokenizer) > model.get_input_embeddings().weight.shape[0]: print("WARNING: Resizing the embedding matrix to match the tokenizer vocab size.") model.resize_token_embeddings(len(tokenizer)) print_model_size(model, train_config, rank if train_config.enable_fsdp else 0) # Prepare the model for int8 training if quantization is enabled if train_config.quantization: model = prepare_model_for_kbit_training(model) # Convert the model to bfloat16 if fsdp and pure_bf16 is enabled if train_config.enable_fsdp and fsdp_config.pure_bf16: model.to(torch.bfloat16) if train_config.use_peft: peft_config = generate_peft_config(train_config, kwargs) model = get_peft_model(model, peft_config) model.print_trainable_parameters() if wandb_run: wandb_run.config.update(peft_config) hsdp_device_mesh = None if fsdp_config.hsdp and fsdp_config.sharding_strategy == ShardingStrategy.HYBRID_SHARD: hsdp_device_mesh = hsdp_device_mesh(replica_group_size=fsdp_config.replica_group_size, sharding_group_size=fsdp_config.sharding_group_size) print("HSDP device mesh is ready") #setting up FSDP if enable_fsdp is enabled if train_config.enable_fsdp: if not train_config.use_peft and train_config.freeze_layers: freeze_transformer_layers(train_config.num_freeze_layers) mixed_precision_policy, wrapping_policy = get_policies(fsdp_config, rank) my_auto_wrapping_policy = fsdp_auto_wrap_policy(model, LlamaDecoderLayer) device_id = 0 if is_xpu_available(): device_id = torch.xpu.current_device() elif torch.cuda.is_available(): device_id = torch.cuda.current_device() model = FSDP( model, auto_wrap_policy= my_auto_wrapping_policy if train_config.use_peft else wrapping_policy, cpu_offload=CPUOffload(offload_params=True) if fsdp_config.fsdp_cpu_offload else None, mixed_precision=mixed_precision_policy if not fsdp_config.pure_bf16 else None, sharding_strategy=fsdp_config.sharding_strategy, device_mesh=hsdp_device_mesh, device_id=device_id, limit_all_gathers=True, sync_module_states=train_config.low_cpu_fsdp, param_init_fn=lambda module: module.to_empty(device=torch.device("cuda"), recurse=False) if train_config.low_cpu_fsdp and rank != 0 else None, ) if fsdp_config.fsdp_activation_checkpointing: apply_fsdp_checkpointing(model) elif not train_config.quantization and not train_config.enable_fsdp: if is_xpu_available(): model.to("xpu:0") elif torch.cuda.is_available(): model.to("cuda") dataset_config = generate_dataset_config(train_config, kwargs) # Load and preprocess the dataset for training and validation dataset_train = get_preprocessed_dataset( tokenizer, dataset_config, split="train", ) if not train_config.enable_fsdp or rank == 0: print(f"--> Training Set Length = {len(dataset_train)}") dataset_val = get_preprocessed_dataset( tokenizer, dataset_config, split="test", ) if not train_config.enable_fsdp or rank == 0: print(f"--> Validation Set Length = {len(dataset_val)}") if train_config.batching_strategy == "packing": dataset_train = ConcatDataset(dataset_train, chunk_size=train_config.context_length) train_dl_kwargs = get_dataloader_kwargs(train_config, dataset_train, tokenizer, "train") # Create DataLoaders for the training and validation dataset train_dataloader = torch.utils.data.DataLoader( dataset_train, num_workers=train_config.num_workers_dataloader, pin_memory=True, **train_dl_kwargs, ) eval_dataloader = None if train_config.run_validation: if train_config.batching_strategy == "packing": dataset_val = ConcatDataset(dataset_val, chunk_size=train_config.context_length) val_dl_kwargs = get_dataloader_kwargs(train_config, dataset_val, tokenizer, "val") eval_dataloader = torch.utils.data.DataLoader( dataset_val, num_workers=train_config.num_workers_dataloader, pin_memory=True, **val_dl_kwargs, ) # Initialize the optimizer and learning rate scheduler if fsdp_config.pure_bf16 and fsdp_config.optimizer == "anyprecision": optimizer = AnyPrecisionAdamW( model.parameters(), lr=train_config.lr, momentum_dtype=torch.bfloat16, variance_dtype=torch.bfloat16, use_kahan_summation=False, weight_decay=train_config.weight_decay, ) else: optimizer = optim.AdamW( model.parameters(), lr=train_config.lr, weight_decay=train_config.weight_decay, ) scheduler = StepLR(optimizer, step_size=1, gamma=train_config.gamma) # Start the training process results = train( model, train_dataloader, eval_dataloader, tokenizer, optimizer, scheduler, train_config.gradient_accumulation_steps, train_config, fsdp_config if train_config.enable_fsdp else None, local_rank if train_config.enable_fsdp else None, rank if train_config.enable_fsdp else None, wandb_run, ) if not train_config.enable_fsdp or rank==0: [print(f'Key: {k}, Value: {v}') for k, v in results.items()] if train_config.use_wandb: for k,v in results.items(): wandb_run.summary[k] = v if __name__ == "__main__": fire.Fire(main)