github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[feat] support meta cache, meta_grad_send, meta_tensor_send; fix runtime too long in Recv Bwd; benchmark for llama + Hybrid(tp+pp);

pull/6083/head
duanjunwen 2024-10-24 07:30:19 +00:00
parent 705b18e1e7
commit 2eca112c90
8 changed files with 184 additions and 63 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

107
colossalai/pipeline/schedule/zero_bubble_pp.py
View File

@ -8,7 +8,7 @@ from torch.utils._pytree import tree_flatten, tree_map
from colossalai.accelerator import get_accelerator from colossalai.accelerator import get_accelerator
from colossalai.interface import OptimizerWrapper from colossalai.interface import OptimizerWrapper
from colossalai.pipeline.p2p import PipelineP2PCommunication from colossalai.pipeline.p2p import PipelineP2PCommunication, create_send_metadata
from colossalai.pipeline.schedule.v_schedule import ScheduledNode from colossalai.pipeline.schedule.v_schedule import ScheduledNode
from colossalai.pipeline.stage_manager import PipelineStageManager from colossalai.pipeline.stage_manager import PipelineStageManager
from colossalai.pipeline.weight_grad_store import WeightGradStore from colossalai.pipeline.weight_grad_store import WeightGradStore
@ -62,11 +62,11 @@ class ZeroBubbleVPipeScheduler(PipelineSchedule):
self.do_post_validation = False self.do_post_validation = False
# P2PMeta cache # P2PMeta cache
# self.enable_metadata_cache = enable_metadata_cache self.enable_metadata_cache = enable_metadata_cache
# self.send_tensor_metadata = True self.send_tensor_metadata = True
# self.send_grad_metadata = True self.send_grad_metadata = True
# self.tensor_metadata_recv = None self.tensor_metadata_recv = None
# self.grad_metadata_recv = None self.grad_metadata_recv = None
# P2P communication # P2P communication
self.comm = PipelineP2PCommunication(stage_manager, overlap_p2p=overlap_p2p) self.comm = PipelineP2PCommunication(stage_manager, overlap_p2p=overlap_p2p)
@ -105,8 +105,11 @@ class ZeroBubbleVPipeScheduler(PipelineSchedule):
# dy buffer for local send bwd # dy buffer for local send bwd
self.local_send_backward_buffer = [] self.local_send_backward_buffer = []
# wait pp buffer
self.send_handles = []
def assert_buffer_empty(self): def assert_buffer_empty(self):
# assert buuffer is empty at end # assert buffer is empty at end
assert len(self.input_tensors[0]) == 0 assert len(self.input_tensors[0]) == 0
assert len(self.input_tensors[1]) == 0 assert len(self.input_tensors[1]) == 0
assert len(self.output_tensors[0]) == 0 assert len(self.output_tensors[0]) == 0
@ -125,6 +128,7 @@ class ZeroBubbleVPipeScheduler(PipelineSchedule):
assert len(self.recv_backward_buffer[1]) == 0 assert len(self.recv_backward_buffer[1]) == 0
assert len(self.local_send_forward_buffer) == 0 assert len(self.local_send_forward_buffer) == 0
assert len(self.local_send_backward_buffer) == 0 assert len(self.local_send_backward_buffer) == 0
# assert len(self.send_handles) == 0
def load_batch(self, data_iter: Iterable, device: Optional[torch.device] = None) -> None: def load_batch(self, data_iter: Iterable, device: Optional[torch.device] = None) -> None:
"""Load a batch from data iterator. """Load a batch from data iterator.
@ -221,7 +225,8 @@ class ZeroBubbleVPipeScheduler(PipelineSchedule):
# do nothing; cause u are chunk 0 in first rank, u have no prev rank; # do nothing; cause u are chunk 0 in first rank, u have no prev rank;
################# #################
if self.stage_manager.is_first_stage(ignore_chunk=True): if self.stage_manager.is_first_stage(ignore_chunk=True):
return None, [] # return None, []
return []
################ ################
# chunk = 0 & not is_first_stage # chunk = 0 & not is_first_stage
@ -229,9 +234,14 @@ class ZeroBubbleVPipeScheduler(PipelineSchedule):
################# #################
else: else:
prev_rank = self.stage_manager.get_prev_rank() prev_rank = self.stage_manager.get_prev_rank()
input_tensor, wait_handles = self.comm.recv_forward(prev_rank=prev_rank) input_tensor, wait_handles = self.comm.recv_forward(
prev_rank=prev_rank, metadata_recv=self.tensor_metadata_recv
)
if self.enable_metadata_cache and self.tensor_metadata_recv is None:
self.tensor_metadata_recv = create_send_metadata(input_tensor)
self.recv_forward_buffer[model_chunk_id].append(input_tensor) self.recv_forward_buffer[model_chunk_id].append(input_tensor)
return input_tensor, wait_handles # return input_tensor, wait_handles
return wait_handles
else: else:
################ ################
@ -239,7 +249,8 @@ class ZeroBubbleVPipeScheduler(PipelineSchedule):
# do nothing; cause u get y from local_send_forward_buffer in schedule f # do nothing; cause u get y from local_send_forward_buffer in schedule f
################ ################
if self.stage_manager.is_last_stage(ignore_chunk=True): if self.stage_manager.is_last_stage(ignore_chunk=True):
return None, [] # return None, []
return []
################ ################
# chunk = 1 & not is_last_stage # chunk = 1 & not is_last_stage
@ -247,9 +258,14 @@ class ZeroBubbleVPipeScheduler(PipelineSchedule):
################ ################
else: else:
next_rank = self.stage_manager.get_next_rank() next_rank = self.stage_manager.get_next_rank()
input_tensor, wait_handles = self.comm.recv_forward(next_rank) input_tensor, wait_handles = self.comm.recv_forward(
next_rank, metadata_recv=self.tensor_metadata_recv
)
if self.enable_metadata_cache and self.tensor_metadata_recv is None:
self.tensor_metadata_recv = create_send_metadata(input_tensor)
self.recv_forward_buffer[model_chunk_id].append(input_tensor) self.recv_forward_buffer[model_chunk_id].append(input_tensor)
return input_tensor, wait_handles # return input_tensor, wait_handles
return wait_handles
def recv_backward(self, model_chunk_id: int, next_rank: int = None) -> Tuple[Any, List]: def recv_backward(self, model_chunk_id: int, next_rank: int = None) -> Tuple[Any, List]:
"""Copy the gradient tensor from the next stage in pipeline as the input gradient of this stage. """Copy the gradient tensor from the next stage in pipeline as the input gradient of this stage.
@ -271,7 +287,8 @@ class ZeroBubbleVPipeScheduler(PipelineSchedule):
# do nothing; Already get dy from local_send_backward_buffer in schedule b # do nothing; Already get dy from local_send_backward_buffer in schedule b
################ ################
if self.stage_manager.is_last_stage(ignore_chunk=True): if self.stage_manager.is_last_stage(ignore_chunk=True):
return None, [] # return None, []
return []
################ ################
# chunk = 0 & not is_last_stage # chunk = 0 & not is_last_stage
@ -279,9 +296,14 @@ class ZeroBubbleVPipeScheduler(PipelineSchedule):
################ ################
else: else:
next_rank = self.stage_manager.get_next_rank() next_rank = self.stage_manager.get_next_rank()
output_tensor_grad, wait_handles = self.comm.recv_backward(next_rank) output_tensor_grad, wait_handles = self.comm.recv_backward(
next_rank, metadata_recv=self.grad_metadata_recv
)
if self.enable_metadata_cache and self.grad_metadata_recv is None:
self.grad_metadata_recv = create_send_metadata(output_tensor_grad)
self.recv_backward_buffer[model_chunk_id].append(output_tensor_grad) self.recv_backward_buffer[model_chunk_id].append(output_tensor_grad)
return output_tensor_grad, wait_handles # return output_tensor_grad, wait_handles
return wait_handles
else: else:
# bwd chunk1 is left V; # bwd chunk1 is left V;
@ -290,7 +312,8 @@ class ZeroBubbleVPipeScheduler(PipelineSchedule):
# do nothing; get loss from local # do nothing; get loss from local
################ ################
if self.stage_manager.is_first_stage(ignore_chunk=True): if self.stage_manager.is_first_stage(ignore_chunk=True):
return None, [] # return None, []
return []
################ ################
# chunk = 1 & not first stage # chunk = 1 & not first stage
@ -298,9 +321,14 @@ class ZeroBubbleVPipeScheduler(PipelineSchedule):
################ ################
else: else:
prev_rank = self.stage_manager.get_prev_rank() prev_rank = self.stage_manager.get_prev_rank()
output_tensor_grad, wait_handles = self.comm.recv_backward(next_rank=prev_rank) output_tensor_grad, wait_handles = self.comm.recv_backward(
next_rank=prev_rank, metadata_recv=self.grad_metadata_recv
)
if self.enable_metadata_cache and self.grad_metadata_recv is None:
self.grad_metadata_recv = create_send_metadata(output_tensor_grad)
self.recv_backward_buffer[model_chunk_id].append(output_tensor_grad) self.recv_backward_buffer[model_chunk_id].append(output_tensor_grad)
return output_tensor_grad, wait_handles # return output_tensor_grad, wait_handles
return wait_handles
def send_forward(self, model_chunk_id: int, next_rank: int = None) -> List: def send_forward(self, model_chunk_id: int, next_rank: int = None) -> List:
"""Sends the input tensor to the next stage in pipeline. """Sends the input tensor to the next stage in pipeline.
@ -330,7 +358,10 @@ class ZeroBubbleVPipeScheduler(PipelineSchedule):
else: else:
next_rank = self.stage_manager.get_next_rank() next_rank = self.stage_manager.get_next_rank()
output_tensor = self.send_forward_buffer[model_chunk_id].pop(0) output_tensor = self.send_forward_buffer[model_chunk_id].pop(0)
send_handles = self.comm.send_forward(output_object=output_tensor, next_rank=next_rank) send_handles = self.comm.send_forward(
output_object=output_tensor, next_rank=next_rank, send_metadata=self.send_tensor_metadata
)
self.send_tensor_metadata = not self.enable_metadata_cache
return send_handles return send_handles
else: else:
@ -348,7 +379,10 @@ class ZeroBubbleVPipeScheduler(PipelineSchedule):
else: else:
prev_rank = self.stage_manager.get_prev_rank() prev_rank = self.stage_manager.get_prev_rank()
output_tensor = self.send_forward_buffer[model_chunk_id].pop(0) output_tensor = self.send_forward_buffer[model_chunk_id].pop(0)
send_handles = self.comm.send_forward(output_tensor, prev_rank) send_handles = self.comm.send_forward(
output_tensor, prev_rank, send_metadata=self.send_tensor_metadata
)
self.send_tensor_metadata = not self.enable_metadata_cache
return send_handles return send_handles
def send_backward(self, model_chunk_id: int, prev_rank: int = None) -> List: def send_backward(self, model_chunk_id: int, prev_rank: int = None) -> List:
@ -380,7 +414,10 @@ class ZeroBubbleVPipeScheduler(PipelineSchedule):
else: else:
prev_rank = self.stage_manager.get_prev_rank() prev_rank = self.stage_manager.get_prev_rank()
input_tensor_grad = self.send_backward_buffer[model_chunk_id].pop(0) input_tensor_grad = self.send_backward_buffer[model_chunk_id].pop(0)
send_handles = self.comm.send_backward(input_tensor_grad, prev_rank) send_handles = self.comm.send_backward(
input_tensor_grad, prev_rank, send_metadata=self.send_grad_metadata
)
self.send_grad_metadata = not self.enable_metadata_cache
return send_handles return send_handles
# bwd chunk1 is left V; # bwd chunk1 is left V;
@ -399,7 +436,10 @@ class ZeroBubbleVPipeScheduler(PipelineSchedule):
else: else:
next_rank = self.stage_manager.get_next_rank() next_rank = self.stage_manager.get_next_rank()
input_tensor_grad = self.send_backward_buffer[model_chunk_id].pop(0) input_tensor_grad = self.send_backward_buffer[model_chunk_id].pop(0)
send_handles = self.comm.send_backward(input_tensor_grad, next_rank) send_handles = self.comm.send_backward(
input_tensor_grad, next_rank, send_metadata=self.send_grad_metadata
)
self.send_grad_metadata = not self.enable_metadata_cache
return send_handles return send_handles
def forward_step( def forward_step(
@ -479,11 +519,11 @@ class ZeroBubbleVPipeScheduler(PipelineSchedule):
output_obj_grad_ = [] output_obj_grad_ = []
# For chunk 0 stage 0, use micro_batch as input_obj_; and we don't have to cal microbatch dx. # For chunk 0 stage 0, use micro_batch as input_obj_; and we don't have to cal microbatch dx.
if model_chunk_id == 0 and self.stage_manager.is_first_stage(ignore_chunk=True): # if model_chunk_id == 0 and self.stage_manager.is_first_stage(ignore_chunk=True):
return None # return None
# For loss backward; output_obj is loss; output_obj_grad should be None # For loss backward; output_obj is loss; output_obj_grad should be None
elif model_chunk_id == 1 and self.stage_manager.is_first_stage(ignore_chunk=True): if model_chunk_id == 1 and self.stage_manager.is_first_stage(ignore_chunk=True):
assert output_obj_grad is None assert output_obj_grad is None
input_obj_, _ = tree_flatten(input_obj) input_obj_, _ = tree_flatten(input_obj)
output_obj_.append(output_obj) # LOSS output_obj_.append(output_obj) # LOSS
@ -510,7 +550,7 @@ class ZeroBubbleVPipeScheduler(PipelineSchedule):
tensor=output_obj_, tensor=output_obj_,
grad=output_obj_grad_, grad=output_obj_grad_,
# inputs=input_obj_, # inputs=input_obj_,
# retain_graph=True, retain_graph=False,
) )
# Format output_obj_grad # Format output_obj_grad
input_obj_grad = dict() input_obj_grad = dict()
@ -712,6 +752,12 @@ class ZeroBubbleVPipeScheduler(PipelineSchedule):
# else: # else:
# # we save output_tensor_grad here # # we save output_tensor_grad here
# self.output_tensors_grad_dw[model_chunk_id].append(output_tensor_grad) # self.output_tensors_grad_dw[model_chunk_id].append(output_tensor_grad)
# the_output_obj_grad = []
# if isinstance(output_obj, dict):
# for (k, v) in output_obj.items():
# the_output_obj_grad.append(v.requires_grad)
# else:
# the_output_obj_grad.append(output_obj.requires_grad)
input_object_grad = self.backward_b_step( input_object_grad = self.backward_b_step(
model_chunk=model_chunk, model_chunk=model_chunk,
@ -844,7 +890,8 @@ class ZeroBubbleVPipeScheduler(PipelineSchedule):
if scheduled_node.type in AUTO_SCHEDULE_COMMUNICATION_TYPES: if scheduled_node.type in AUTO_SCHEDULE_COMMUNICATION_TYPES:
# communication # communication
communication_func = self.communication_map[scheduled_node.type] communication_func = self.communication_map[scheduled_node.type]
communication_func(scheduled_node.chunk) wait_handle = communication_func(scheduled_node.chunk)
self.send_handles.append(wait_handle)
elif scheduled_node.type == "F": elif scheduled_node.type == "F":
self.schedule_f( self.schedule_f(
scheduled_node=scheduled_node, scheduled_node=scheduled_node,
@ -868,6 +915,9 @@ class ZeroBubbleVPipeScheduler(PipelineSchedule):
model_chunk_id=scheduled_node.chunk, model_chunk_id=scheduled_node.chunk,
optimizer=optimizer, optimizer=optimizer,
) )
for h in self.send_handles:
for hh in h:
hh.wait()
# return loss & output # return loss & output
if outputs is not None: if outputs is not None:
@ -907,5 +957,4 @@ class ZeroBubbleVPipeScheduler(PipelineSchedule):
) )
self.assert_buffer_empty() self.assert_buffer_empty()
return result return result

2
colossalai/pipeline/stage_manager.py
View File

@ -223,10 +223,10 @@ class PipelineStageManager:
# calculate the num_layers per stage # calculate the num_layers per stage
layers_per_stage = [quotient] * num_stages * num_model_chunks layers_per_stage = [quotient] * num_stages * num_model_chunks
# deal with the rest layers # deal with the rest layers
if remainder > 0: if remainder > 0:
start_position = (num_stages * num_model_chunks) // 2 - remainder // 2 start_position = (num_stages * num_model_chunks) // 2 - remainder // 2
for i in range(start_position, start_position + remainder): for i in range(start_position, start_position + remainder):
layers_per_stage[i] += 1 layers_per_stage[i] += 1
# print(f"layers_per_stage {layers_per_stage}")
return layers_per_stage return layers_per_stage

55
colossalai/pipeline/weight_grad_store.py
View File

@ -1,9 +1,6 @@
import queue import queue
# from megatron import get_args from colossalai.pipeline.stage_manager import PipelineStageManager
# from megatron.core import parallel_state
# from megatron.core.distributed.finalize_model_grads import _allreduce_embedding_grads
# from megatron.core.utils import get_model_config, get_attr_wrapped_model
class WeightGradStore: class WeightGradStore:
@ -23,6 +20,7 @@ class WeightGradStore:
@classmethod @classmethod
def pop(cls, chunk=0): def pop(cls, chunk=0):
# print(f"chunk id {chunk} queue size {cls.weight_grad_queue[chunk].qsize()}")
if cls.weight_grad_queue[chunk].qsize() > 0: if cls.weight_grad_queue[chunk].qsize() > 0:
stored_grads = cls.weight_grad_queue[chunk].get() stored_grads = cls.weight_grad_queue[chunk].get()
for total_input, grad_output, weight, func in stored_grads: for total_input, grad_output, weight, func in stored_grads:
@ -34,3 +32,52 @@ class WeightGradStore:
weight.grad = grad_weight weight.grad = grad_weight
else: else:
raise Exception("Pop empty queue.") raise Exception("Pop empty queue.")
@classmethod
def clear(cls, stage_manager: PipelineStageManager, chunk=0):
pass
# print(f"stage {stage_manager.stage} len_chunk_0 {cls.weight_grad_queue[0].qsize()} len_chunk_1 {cls.weight_grad_queue[1].qsize()}")
# while cls.weight_grad_queue[chunk].qsize() > 0:
# stored_grads = cls.weight_grad_queue[chunk].get()
# for total_input, grad_output, weight, func in stored_grads:
# if weight.grad is not None:
# func(total_input, grad_output, weight.grad)
# # for first bwd; weight.grad is None, assign grad_weight to weight.grad
# else:
# grad_weight = func(total_input, grad_output)
# weight.grad = grad_weight
# weight_grad_tasks = []
# while cls.weight_grad_queue[chunk].qsize() > 0:
# stored_grads = cls.weight_grad_queue[chunk].get()
# if len(weight_grad_tasks) == 0:
# for _ in stored_grads:
# weight_grad_tasks.append([])
# else:
# assert len(weight_grad_tasks) == len(stored_grads)
# for i, task in enumerate(stored_grads):
# weight_grad_tasks[i].append(task)
# if stage_manager.is_last_stage(ignore_chunk=True) and chunk == 1:
# assert len(weight_grad_tasks) > 0
# output_layer_grads = weight_grad_tasks[0]
# for j in range(len(output_layer_grads)):
# total_input, grad_output, weight, func = output_layer_grads[j]
# if output_layer_weight is None:
# output_layer_weight = weight
# assert output_layer_weight is weight
# func(total_input, grad_output, weight.grad)
# output_layer_grads[j] = None # release memory
# weight_grad_tasks = weight_grad_tasks[1:]
# for i in range(len(weight_grad_tasks)):
# tasks = weight_grad_tasks[i]
# param = None
# for j in range(len(tasks)):
# total_input, grad_output, weight, func = tasks[j]
# if param is None:
# param = weight
# assert param is weight
# func(total_input, grad_output, weight.grad)
# tasks[j] = None # release memory
# weight_grad_tasks[i] = None # release memory

7
colossalai/shardformer/modeling/llama.py
View File

@ -32,6 +32,7 @@ from colossalai.shardformer.shard import ShardConfig
from ..layer import ColoAttention, RingAttention, dist_cross_entropy from ..layer import ColoAttention, RingAttention, dist_cross_entropy
_SUPPORTED_SP_MODE = ["all_to_all", "split_gather", "ring", "ring_attn"] _SUPPORTED_SP_MODE = ["all_to_all", "split_gather", "ring", "ring_attn"]
_GLOBAL_ORDER_ = 0
class LlamaPipelineForwards: class LlamaPipelineForwards:
@ -193,6 +194,10 @@ class LlamaPipelineForwards:
assert num_ckpt_layers <= end_idx - start_idx assert num_ckpt_layers <= end_idx - start_idx
for idx, decoder_layer in enumerate(self.layers[start_idx:end_idx], start=start_idx): for idx, decoder_layer in enumerate(self.layers[start_idx:end_idx], start=start_idx):
# global _GLOBAL_ORDER_
# if torch.distributed.get_rank() == 0:
# print(f"rank {torch.distributed.get_rank()} {stage_manager.stage}; start:{start_idx}, end:{end_idx} hidden_states require grad{hidden_states.requires_grad}")
# # _GLOBAL_ORDER_ += 1
if output_hidden_states: if output_hidden_states:
all_hidden_states += (hidden_states,) all_hidden_states += (hidden_states,)
if idx - start_idx < num_ckpt_layers: if idx - start_idx < num_ckpt_layers:
@ -216,6 +221,8 @@ class LlamaPipelineForwards:
use_cache=use_cache, use_cache=use_cache,
cache_position=cache_position, cache_position=cache_position,
) )
# if torch.distributed.get_rank() == 0:
# print(f"rank {torch.distributed.get_rank()} {stage_manager.stage}; start:{start_idx}, end:{end_idx} layer_outputs require grad {layer_outputs[0].requires_grad}")
hidden_states = layer_outputs[0] hidden_states = layer_outputs[0]
if use_cache: if use_cache:

27
colossalai/shardformer/policies/llama.py
View File

@ -96,7 +96,7 @@ class LlamaPolicy(Policy):
target_key=attn_cls, target_key=attn_cls,
) )
if self.pipeline_stage_manager is None: if self.pipeline_stage_manager is not None:
self.append_or_create_method_replacement( self.append_or_create_method_replacement(
description={ description={
"forward": partial( "forward": partial(
@ -298,7 +298,6 @@ class LlamaPolicy(Policy):
not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True)
): ):
held_layers.append(module.norm) held_layers.append(module.norm)
else: else:
layers_per_stage = stage_manager.distribute_layers(len(module.layers)) layers_per_stage = stage_manager.distribute_layers(len(module.layers))
if stage_manager.is_first_stage(): if stage_manager.is_first_stage():
@ -395,8 +394,8 @@ class LlamaForCausalLMPolicy(LlamaPolicy):
return held_layers return held_layers
def get_shared_params(self) -> List[Dict[int, Tensor]]: def get_shared_params(self) -> List[Dict[int, Tensor]]:
if self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv: # if self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv:
return [] # return []
llama_model = self.model.model llama_model = self.model.model
if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1: if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1:
if ( if (
@ -404,12 +403,20 @@ class LlamaForCausalLMPolicy(LlamaPolicy):
and self.pipeline_stage_manager.num_stages > 1 and self.pipeline_stage_manager.num_stages > 1
): ):
# tie weights # tie weights
return [ if self.pipeline_stage_manager.use_zbv:
{ return [
0: llama_model.embed_tokens.weight, {
self.pipeline_stage_manager.num_stages - 1: self.model.lm_head.weight, 0: llama_model.embed_tokens.weight,
} 0: self.model.lm_head.weight,
] }
]
else:
return [
{
0: llama_model.embed_tokens.weight,
self.pipeline_stage_manager.num_stages - 1: self.model.lm_head.weight,
}
]
return [] return []

20
examples/language/llama/benchmark.py
View File

@ -40,6 +40,7 @@ MODEL_CONFIGS = {
), ),
"5b": LlamaConfig(max_position_embeddings=4096, num_key_value_heads=8), "5b": LlamaConfig(max_position_embeddings=4096, num_key_value_heads=8),
"7b": LlamaConfig(max_position_embeddings=4096), "7b": LlamaConfig(max_position_embeddings=4096),
# "7b": LlamaConfig(num_hidden_layers=4, max_position_embeddings=4096),
"13b": LlamaConfig( "13b": LlamaConfig(
hidden_size=5120, hidden_size=5120,
intermediate_size=13824, intermediate_size=13824,
@ -127,9 +128,12 @@ def main():
{ {
"gradient_checkpoint_config": PipelineGradientCheckpointConfig( "gradient_checkpoint_config": PipelineGradientCheckpointConfig(
num_ckpt_layers_per_stage=[19, 19, 19, 13], num_ckpt_layers_per_stage=[19, 19, 19, 13],
# num_ckpt_layers_per_stage=[48, 48, 48, 48],
), ),
"num_layers_per_stage": [19, 20, 20, 21], "num_layers_per_stage": [19, 20, 20, 21],
"pp_style": "interleaved", # "num_layers_per_stage": [48, 48, 48, 48],
# "pp_style": "interleaved",
"pp_style": "1f1b",
} }
if args.custom_ckpt if args.custom_ckpt
else {} else {}
@ -227,12 +231,14 @@ def main():
b_cost=1000, b_cost=1000,
w_cost=1000, w_cost=1000,
c_cost=1, c_cost=1,
f_mem=mem_f, f_mem=mem_f * 1.5,
b_mem=mem_b, b_mem=mem_b * 1.5,
w_mem=mem_w, w_mem=mem_w * 1.5,
).get_v_schedule() ).get_v_schedule()
else: else:
scheduler_nodes = None scheduler_nodes = None
# print(f"{dist.get_rank()} {scheduler_nodes[]} ")
plugin = HybridParallelPlugin( plugin = HybridParallelPlugin(
tp_size=args.tp, tp_size=args.tp,
pp_size=args.pp, pp_size=args.pp,
@ -267,7 +273,7 @@ def main():
microbatch_size=args.mbs, microbatch_size=args.mbs,
initial_scale=2**8, initial_scale=2**8,
precision="bf16", precision="bf16",
overlap_p2p=args.overlap, overlap_p2p=True,
use_fp8=args.use_fp8, use_fp8=args.use_fp8,
fp8_communication=args.use_fp8_comm, fp8_communication=args.use_fp8_comm,
) )
@ -328,7 +334,7 @@ def main():
torch.set_default_dtype(torch.bfloat16) torch.set_default_dtype(torch.bfloat16)
model, optimizer, _, dataloader, _ = booster.boost(model, optimizer, dataloader=dataloader) model, optimizer, _, dataloader, _ = booster.boost(model, optimizer, dataloader=dataloader)
torch.set_default_dtype(torch.float) # torch.set_default_dtype(torch.float)
coordinator.print_on_master( coordinator.print_on_master(
f"Booster init max CUDA memory: {get_accelerator().max_memory_allocated()/1024**2:.2f} MB" f"Booster init max CUDA memory: {get_accelerator().max_memory_allocated()/1024**2:.2f} MB"
) )
@ -340,7 +346,7 @@ def main():
args.profile, args.profile,
args.ignore_steps, args.ignore_steps,
1, # avoid creating massive log files 1, # avoid creating massive log files
save_dir=f"profile/{time.strftime('%H:%M', time.localtime())}-{args.plugin}-llama-{args.config}", save_dir=f"./profile/{time.strftime('%H:%M', time.localtime())}-{args.plugin}-llama-{args.config}",
nsys=args.nsys, nsys=args.nsys,
) as prof: ) as prof:
if isinstance(plugin, HybridParallelPlugin) and args.pp > 1: if isinstance(plugin, HybridParallelPlugin) and args.pp > 1:

13
examples/language/performance_evaluator.py
View File

@ -21,11 +21,16 @@ def divide(x: float, y: float) -> float:
def all_reduce_mean(x: float, world_size: int) -> float: def all_reduce_mean(x: float, world_size: int) -> float:
if world_size == 1: if world_size == 1:
return x return x
# BUG: RuntimeError: Invalid scalar type when use dist.all_reduce(tensor, group=gloo_group)
# # Use CPU tensor to avoid OOM/weird NCCl error
# gloo_group = dist.new_group(backend="gloo")
# tensor = torch.tensor([x], device="cpu")
# dist.all_reduce(tensor, group=gloo_group)
# tensor = tensor / world_size
# return tensor.item()
# Use CPU tensor to avoid OOM/weird NCCl error tensor = torch.tensor([x], device=torch.cuda.current_device(), dtype=torch.float)
gloo_group = dist.new_group(backend="gloo") dist.all_reduce(tensor)
tensor = torch.tensor([x], device="cpu")
dist.all_reduce(tensor, group=gloo_group)
tensor = tensor / world_size tensor = tensor / world_size
return tensor.item() return tensor.item()

16
tests/test_pipeline/test_schedule/test_zerobubble_pp.py
View File

@ -758,11 +758,11 @@ def run_with_hybridplugin(test_config):
@parameterize( @parameterize(
"config", "config",
[ [
(0, 1, 4, 1, 1), # (0, 1, 4, 1, 1),
(1, 2, 2, 1, 1), # (1, 2, 2, 1, 1),
(1, 1, 2, 2, 1), (1, 1, 2, 2, 1),
(1, 2, 1, 2, 1), # (1, 2, 1, 2, 1),
(1, 2, 1, 1, 2), # (1, 2, 1, 1, 2),
], ],
) )
def run_with_booster_moehybridplugin(config: Tuple[int, ...]): def run_with_booster_moehybridplugin(config: Tuple[int, ...]):
@ -923,10 +923,10 @@ def run_with_booster_moehybridplugin(config: Tuple[int, ...]):
@parameterize( @parameterize(
"config", "config",
[ [
(0, 4, 1, 1), # (0, 4, 1, 1),
(1, 2, 2, 1), (1, 2, 2, 1),
(1, 2, 1, 2), # (1, 2, 1, 2),
(1, 1, 2, 2), # (1, 1, 2, 2), # TODO: no pp show gather result err
], ],
) )
def run_with_booster_hybridplugin(config: Tuple[int, ...]): def run_with_booster_hybridplugin(config: Tuple[int, ...]):
@ -976,7 +976,7 @@ def run_with_booster_hybridplugin(config: Tuple[int, ...]):
zbv_schedule = graph.get_v_schedule() zbv_schedule = graph.get_v_schedule()
# init MoeHybridPlugin # init HybridParallelPlugin
plugin = HybridParallelPlugin( plugin = HybridParallelPlugin(
pp_size=pp_size, pp_size=pp_size,
num_microbatches=pp_size, num_microbatches=pp_size,