[shardformer] update llama2/opt finetune example and fix llama2 policy (

#4645) * [shardformer] update shardformer readme [shardformer] update shardformer readme [shardformer] update shardformer readme * [shardformer] update llama2/opt finetune example and shardformer update to llama2 * [shardformer] update llama2/opt finetune example and shardformer update to llama2 * [shardformer] update llama2/opt finetune example and shardformer update to llama2 * [shardformer] change dataset * [shardformer] change dataset * [shardformer] fix CI * [shardformer] fix * [shardformer] fix * [shardformer] fix * [shardformer] fix * [shardformer] fix [example] update opt example [example] resolve comments fix fix
hpcaitech · Sep 9, 2023 · 7486ed7 · 7486ed7
1 parent a686f9d
commit 7486ed7
Show file tree

Hide file tree

Showing 12 changed files with 166 additions and 168 deletions.
diff --git a/colossalai/shardformer/modeling/llama.py b/colossalai/shardformer/modeling/llama.py
@@ -1,3 +1,4 @@
+import warnings
 from typing import Callable, List, Optional, Tuple
 
 import torch
@@ -392,6 +393,13 @@ def get_llama_flash_attention_forward():
 
  from transformers.models.llama.modeling_llama import LlamaAttention, apply_rotary_pos_emb
 
+ llama_version = 2
+ try:
+ from transformers.models.llama.modeling_llama import repeat_kv
+ except:
+ warnings.warn("using llamav1, llamav1 hasn't repeat_kv function")
+ llama_version = 1
+
  from colossalai.kernel.cuda_native import AttnMaskType, ColoAttention
 
  def forward(
@@ -424,6 +432,11 @@ def forward(
 
  past_key_value = (key_states, value_states) if use_cache else None
 
+ # repeat k/v heads if n_kv_heads < n_heads
+ if llama_version == 2:
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
+
  me_input_shape = (bsz, q_len, self.num_heads, self.head_dim)
  query_states = query_states.transpose(1, 2).contiguous().view(*me_input_shape)
  key_states = key_states.transpose(1, 2).contiguous().view(*me_input_shape)

diff --git a/colossalai/shardformer/modeling/opt.py b/colossalai/shardformer/modeling/opt.py
@@ -518,7 +518,6 @@ def forward(
  # for the decoder
  is_cross_attention = key_value_states is not None
  bsz, tgt_len, _ = hidden_states.size()
- assert tgt_len % 4 == 0, "Flash Attention Error: The sequence length should be a multiple of 4."
 
  attention_input_shape = (bsz, -1, self.num_heads, self.head_dim)
  # get query proj

diff --git a/colossalai/shardformer/policies/llama.py b/colossalai/shardformer/policies/llama.py
@@ -43,10 +43,8 @@ def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]:
 
  if self.shard_config.enable_tensor_parallelism:
  decoder_attribute_replacement = {
- "self_attn.hidden_size":
- self.model.config.hidden_size // self.shard_config.tensor_parallel_size,
- "self_attn.num_heads":
- self.model.config.num_attention_heads // self.shard_config.tensor_parallel_size,
+ "self_attn.hidden_size": self.model.config.hidden_size // self.shard_config.tensor_parallel_size,
+ "self_attn.num_heads": self.model.config.num_attention_heads // self.shard_config.tensor_parallel_size,
  }
  if getattr(self.model.config, "num_key_value_heads", False):
  decoder_attribute_replacement["self_attn.num_key_value_heads"] = \

diff --git a/examples/language/bert/finetune.py b/examples/language/bert/finetune.py
@@ -58,45 +58,40 @@ def evaluate_model(
  model.eval()
 
  def evaluate_subset(dataloader: DataLoader):
+ use_pipeline = isinstance(booster.plugin, HybridParallelPlugin) and booster.plugin.pp_size > 1
+ is_pp_last_stage = use_pipeline and booster.plugin.stage_manager.is_last_stage()
+
  accum_loss = torch.zeros(1, device=get_current_device())
  for batch in dataloader:
  batch = move_to_cuda(batch)
  labels = batch["labels"]
- batch_size = batch["input_ids"].shape[0]
- if hasattr(booster.plugin, "stage_manager") and booster.plugin.stage_manager is not None:
+ if use_pipeline:
  pg_mesh = booster.plugin.pg_mesh
  pp_group = booster.plugin.pp_group
  current_pp_group_ranks = pg_mesh.get_ranks_in_group(pp_group)
  current_rank = dist.get_rank()
- #TODO pass dataloader to execute_pipeline directly
  batch = iter([batch])
  outputs = booster.execute_pipeline(batch, model, criterion, return_loss=True, return_outputs=True)
 
- if booster.plugin.stage_manager.is_last_stage():
- val_loss = outputs["loss"]
-
+ if is_pp_last_stage:
  logits = outputs["outputs"]["logits"]
-
+ val_loss = outputs["loss"]
  accum_loss.add_(val_loss)
 
  if num_labels > 1:
  preds = torch.argmax(logits, axis=1)
  elif num_labels == 1:
  preds = logits.squeeze()
 
- dist.broadcast(preds, src=current_rank, group=pp_group)
- dist.broadcast(val_loss, src=current_rank, group=pp_group)
+ dist.broadcast_object_list([preds, val_loss], src=current_pp_group_ranks[-1], group=pp_group)
 
  metric.add_batch(predictions=preds, references=labels)
  elif current_rank in current_pp_group_ranks:
- val_loss = torch.empty((1,), device=get_current_device())
- preds = torch.empty((batch_size,), dtype=torch.int64, device=get_current_device())
-
- dist.broadcast(preds, src=current_pp_group_ranks[-1], group=pp_group)
- dist.broadcast(val_loss, src=current_pp_group_ranks[-1], group=pp_group)
+ object_list = [None, None]
+ dist.broadcast_object_list(object_list, src=current_pp_group_ranks[-1], group=pp_group)
 
- accum_loss.add_(val_loss)
- metric.add_batch(predictions=preds, references=labels)
+ metric.add_batch(predictions=object_list[0].to(get_current_device()), references=labels)
+ accum_loss.add_(object_list[1].to(get_current_device()))
 
  else:
  batch = move_to_cuda(batch)
@@ -132,31 +127,33 @@ def evaluate_subset(dataloader: DataLoader):
 def train_epoch(epoch: int, model: nn.Module, optimizer: Optimizer, _criterion: Callable, lr_scheduler: LRScheduler,
  train_dataloader: DataLoader, booster: Booster, coordinator: DistCoordinator):
 
+ use_pipeline = isinstance(booster.plugin, HybridParallelPlugin) and booster.plugin.pp_size > 1
+ is_pp_last_stage = use_pipeline and booster.plugin.stage_manager.is_last_stage()
+ total_step = len(train_dataloader)
+
  model.train()
- is_pp_last_stage = hasattr(
- booster.plugin,
- "stage_manager") and booster.plugin.stage_manager is not None and booster.plugin.stage_manager.is_last_stage()
- with tqdm(train_dataloader,
+ optimizer.zero_grad()
+ train_dataloader_iter = iter(train_dataloader)
+ with tqdm(range(total_step),
  desc=f'Epoch [{epoch + 1}/{NUM_EPOCHS}]',
  disable=not (coordinator.is_master() or is_pp_last_stage)) as pbar:
- for batch in pbar:
- # Forward pass
- batch = move_to_cuda(batch)
- if hasattr(booster.plugin, "stage_manager") and booster.plugin.stage_manager is not None:
- #TODO pass train_dataloader to execute_pipeline directly
- batch = iter([batch])
- outputs = booster.execute_pipeline(batch,
+ # Forward pass
+ for _ in pbar:
+ if use_pipeline:
+ outputs = booster.execute_pipeline(train_dataloader_iter,
  model,
  _criterion,
  optimizer,
  return_loss=True,
  return_outputs=True)
  # Backward and optimize
- if booster.plugin.stage_manager.is_last_stage():
+ if is_pp_last_stage:
  loss = outputs['loss']
  pbar.set_postfix({'loss': loss.item()})
  else:
- outputs = model(**batch)
+ data = next(train_dataloader_iter)
+ data = move_to_cuda(data)
+ outputs = model(**data)
  loss = _criterion(outputs, None)
  # Backward
  booster.backward(loss, optimizer)

diff --git a/examples/language/opt/args.py b/examples/language/opt/args.py
@@ -4,117 +4,65 @@
 def parse_demo_args():
 
  parser = get_default_parser()
- parser.add_argument(
- "--model_name_or_path",
- type=str,
- default="facebook/opt-350m",
- help="Path to pretrained model or model identifier from huggingface.co/models."
- )
- parser.add_argument(
- "--output_path",
- type=str,
- default="./output_model.bin",
- help="The path of your saved model after finetuning."
- )
+ parser.add_argument("--model_name_or_path",
+ type=str,
+ default="facebook/opt-350m",
+ help="Path to pretrained model or model identifier from huggingface.co/models.")
+ parser.add_argument("--output_path",
+ type=str,
+ default="./output_model.bin",
+ help="The path of your saved model after finetuning.")
  parser.add_argument(
  "--plugin",
  type=str,
  default="gemini",
- help="Plugin to use. Valid plugins include 'torch_ddp','torch_ddp_fp16','gemini','low_level_zero'."
- )
- parser.add_argument(
- "--num_epoch",
- type=int,
- default=10,
- help="Number of epochs."
- )
- parser.add_argument(
- "--batch_size",
- type=int,
- default=32,
- help="Batch size (per dp group) for the training dataloader."
- )
- parser.add_argument(
- "--learning_rate",
- type=float,
- default=5e-5,
- help="Initial learning rate (after the potential warmup period) to use."
- )
- parser.add_argument(
- "--warmup_ratio",
- type=float,
- default=0.1,
- help="Ratio of warmup steps against total training steps."
- )
- parser.add_argument(
- "--weight_decay", 
- type=float, 
- default=0.01, 
- help="Weight decay to use."
- )
- parser.add_argument(
- "--seed", 
- type=int, 
- default=42, 
- help="A seed for reproducible training."
- )
+ help=
+ "Plugin to use. Valid plugins include 'torch_ddp','torch_ddp_fp16','gemini','low_level_zero', 'hybrid_parallel'."
+ )
+ parser.add_argument("--num_epoch", type=int, default=10, help="Number of epochs.")
+ parser.add_argument("--batch_size",
+ type=int,
+ default=32,
+ help="Batch size (per dp group) for the training dataloader.")
+ parser.add_argument("--learning_rate",
+ type=float,
+ default=5e-5,
+ help="Initial learning rate (after the potential warmup period) to use.")
+ parser.add_argument("--warmup_ratio",
+ type=float,
+ default=0.1,
+ help="Ratio of warmup steps against total training steps.")
+ parser.add_argument("--weight_decay", type=float, default=0.01, help="Weight decay to use.")
+ parser.add_argument("--seed", type=int, default=42, help="A seed for reproducible training.")
 
  args = parser.parse_args()
  return args
 
 
-
 def parse_benchmark_args():
 
  parser = get_default_parser()
- parser.add_argument(
- "--model_name_or_path",
- type=str,
- default="facebook/opt-125m",
- help="Path to pretrained model or model identifier from huggingface.co/models."
- )
+ parser.add_argument("--model_name_or_path",
+ type=str,
+ default="facebook/opt-125m",
+ help="Path to pretrained model or model identifier from huggingface.co/models.")
  parser.add_argument(
  "--plugin",
  type=str,
  default="gemini",
- help="Plugin to use. Valid plugins include 'torch_ddp','torch_ddp_fp16','gemini','low_level_zero'."
- )
- parser.add_argument(
- "--batch_size",
- type=int,
- default=32,
- help="Batch size (per dp group) for the training dataloader."
- )
- parser.add_argument(
- "--learning_rate",
- type=float,
- default=5e-5,
- help="Initial learning rate (after the potential warmup period) to use."
- )
- parser.add_argument(
- "--weight_decay", 
- type=float, 
- default=0.0, 
- help="Weight decay to use."
- )
- parser.add_argument(
- "--max_train_steps",
- type=int,
- default=20,
- help="Total number of training steps to perform."
- )
- parser.add_argument(
- "--seed", 
- type=int, 
- default=42, 
- help="A seed for reproducible training."
- )
- parser.add_argument(
- "--mem_cap", 
- type=int, 
- default=0, 
- help="Limit on the usage of space for each GPU (in GB)."
- )
+ help="Plugin to use. Valid plugins include 'torch_ddp','torch_ddp_fp16','gemini','low_level_zero'.")
+ parser.add_argument("--batch_size",
+ type=int,
+ default=32,
+ help="Batch size (per dp group) for the training dataloader.")
+ parser.add_argument("--learning_rate",
+ type=float,
+ default=5e-5,
+ help="Initial learning rate (after the potential warmup period) to use.")
+ parser.add_argument("--weight_decay", type=float, default=0.0, help="Weight decay to use.")
+ parser.add_argument("--max_train_steps", type=int, default=20, help="Total number of training steps to perform.")
+ parser.add_argument("--seed", type=int, default=42, help="A seed for reproducible training.")
+ parser.add_argument("--mem_cap", type=int, default=0, help="Limit on the usage of space for each GPU (in GB).")
  args = parser.parse_args()
 
- return args
+ return args