megatron_patch/model/llama2/transformer.py [39:391]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - try: from einops import rearrange except ImportError: rearrange = None try: from flash_attn.flash_attn_interface import flash_attn_unpadded_func except ImportError: try: from flash_attn.flash_attn_interface import flash_attn_varlen_func as flash_attn_unpadded_func except ImportError: flash_attn_unpadded_func = None """ We use the following notation throughout this file: h: hidden size n: number of attention heads p: number of model parallel partitions np: n/p hp: h/p hn: h/n b: batch size s: sequence length l: number of layers Transformer takes input of size [s, b, h] and returns a tensor of the same size. We use the following arguments: hyperparameters: transformer hyperparameters """ class DropPath(MegatronModule): """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). """ def __init__(self, drop_prob=0.): super(DropPath, self).__init__() self.drop_prob = drop_prob def forward(self, hidden_state): if self.drop_prob == 0. or not self.training: return hidden_state keep_prob = 1 - self.drop_prob # work with diff dim tensors, not just 2D ConvNets # hidden_state: [s, b, h] shape = (1,) + (hidden_state.shape[1],) + (1,) * (hidden_state.ndim - 2) random_tensor = keep_prob + \ torch.rand(shape, dtype=hidden_state.dtype, device=hidden_state.device) random_tensor.floor_() # binarize output = hidden_state.div(keep_prob) * random_tensor return output class ParallelMLP(MegatronModule): """MLP. MLP will take the input with h hidden state, project it to 4*h hidden dimension, perform nonlinear transformation, and project the state back into h hidden dimension. """ def __init__(self, config, is_expert=False): super(ParallelMLP, self).__init__() args = get_args() self.add_bias = config.add_bias_linear ffn_hidden_size = config.ffn_hidden_size if config.gated_linear_unit: ffn_hidden_size *= 2 # Project to 4h. If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf self.dense_h_to_4h = tensor_parallel.ColumnParallelLinear( config.hidden_size, ffn_hidden_size, config=config, init_method=config.init_method, bias=self.add_bias, gather_output=False, skip_bias_add=True, is_expert=is_expert, ) self.bias_gelu_fusion = False self.activation_func = None self.swiglu = args.swiglu if args.openai_gelu: self.activation_func = openai_gelu elif args.onnx_safe: self.activation_func = erf_gelu elif args.swiglu: def swiglu(x): x = torch.chunk(x, 2, dim=-1) return F.silu(x[0]) * x[1] self.activation_func = swiglu elif args.squared_relu: def squared_relu(x): return torch.pow(F.relu(x), 2) self.activation_func = squared_relu else: self.bias_gelu_fusion = args.bias_gelu_fusion self.activation_func = F.gelu # Project back to h. self.dense_4h_to_h = tensor_parallel.RowParallelLinear( config.ffn_hidden_size, config.hidden_size, config=config, init_method=config.output_layer_init_method, bias=self.add_bias, input_is_parallel=True, skip_bias_add=True, is_expert=is_expert, ) def forward(self, hidden_states): # [s, b, 4hp] intermediate_parallel, bias_parallel = self.dense_h_to_4h(hidden_states) if self.bias_gelu_fusion: assert self.add_bias is True assert self.activation_func == F.gelu intermediate_parallel = bias_gelu_impl(intermediate_parallel, bias_parallel) else: if bias_parallel is not None: intermediate_parallel = intermediate_parallel + bias_parallel intermediate_parallel = self.activation_func(intermediate_parallel) # [s, b, h] output, output_bias = self.dense_4h_to_h(intermediate_parallel) return output, output_bias def sinkhorn(cost, tol=0.0001): cost = torch.exp(cost) d0 = torch.ones(cost.size(0), device=cost.device, dtype=cost.dtype) d1 = torch.ones(cost.size(1), device=cost.device, dtype=cost.dtype) eps = 0.00000001 error = 1e9 d1_old = d1 while error > tol: d0 = (1/d0.size(0))*1/(torch.sum(d1*cost,1) + eps) d1 = (1/d1.size(0))*1/(torch.sum(d0.unsqueeze(1)*cost,0)+eps) error = torch.mean(torch.abs(d1_old-d1)) d1_old = d1 return d1*cost*d0.unsqueeze(1) class SwitchMLP(MegatronModule): """ Routes input to one of N MLP "experts" """ def __init__(self, config): super(SwitchMLP, self).__init__() args = get_args() self.router = torch.nn.Linear(args.hidden_size, args.num_experts) self.expert_parallel = config.expert_parallel self.sequence_parallel = config.sequence_parallel self.add_bias = config.add_bias_linear if self.expert_parallel: assert args.num_experts % mpu.get_data_parallel_world_size() == 0 self.num_local_experts = args.num_experts // mpu.get_data_parallel_world_size() local_expert_indices_offset = mpu.get_data_parallel_rank() * self.num_local_experts self.local_expert_indices = [local_expert_indices_offset + i for i in range(self.num_local_experts)] else: self.num_local_experts = args.num_experts self.local_expert_indices = [i for i in range(self.num_local_experts)] self.local_experts = torch.nn.ModuleList() for i in range(self.num_local_experts): self.local_experts.append(ParallelMLP(config, is_expert=True)) def gather_indices(self, local_indices): """ Gather tensors and concatinate along the first dimension.""" if self.expert_parallel: group = get_tensor_and_data_parallel_group() else: group = get_tensor_model_parallel_group() world_size = torch.distributed.get_world_size(group=group) # Bypass the function if we are using only 1 GPU. if world_size == 1: return local_indices dim_size = list(local_indices.size()) dim_size[0] = dim_size[0] * world_size # TODO pre allocate memory output = torch.empty(dim_size, dtype=local_indices.dtype, device=torch.cuda.current_device()) torch.distributed._all_gather_base( output, local_indices.contiguous(), group=group ) return output def forward(self, hidden_states): # hidden_states: [b, s, h] args = get_args() s = hidden_states.size(0) b = hidden_states.size(1) h = hidden_states.size(2) route = self.router(hidden_states).view(-1, args.num_experts) # TODO (rprenger) Right now we're just using the sinkhorn algorithm # for load balancing. There should be an option to do no load balancing # and the algorithm and parametets should be further tested if self.training: with torch.no_grad(): sinkroute = sinkhorn(route.detach().to(dtype=torch.float32)) _, max_ind = torch.max(sinkroute, dim=1) route = torch.sigmoid(route) max_prob = route[torch.arange(route.size(0)), max_ind] else: route = torch.sigmoid(route) max_prob, max_ind = torch.max(route, dim=1) max_prob = torch.unsqueeze(max_prob, 1) hidden_states = hidden_states.view(-1, hidden_states.size(2)) # TODO (rprenger) TODO this could be made easier to read # Converting [s, b, h] to [s*b, h]. # Each vector could be routed differently if self.sequence_parallel or self.expert_parallel: global_hidden_states = \ gather_from_sequence_parallel_region_to_moe( hidden_states, expert_parallel=self.expert_parallel ) global_indices = self.gather_indices(max_ind) else: global_hidden_states = hidden_states global_indices = max_ind output_total = torch.zeros_like(global_hidden_states) if self.add_bias: output_bias_total = torch.zeros_like(global_hidden_states) for expert_num, expert in enumerate(self.local_experts): local_expert_index = self.local_expert_indices[expert_num] local_indices = (global_indices == local_expert_index).nonzero() hidden = global_hidden_states[local_indices, :] output, output_bias = expert(hidden) output_total[local_indices, :] = output if self.add_bias: output_bias = output_bias.expand_as(output) output_bias_total[local_indices, :] = output_bias if self.sequence_parallel or self.expert_parallel: output_total = \ reduce_scatter_to_sequence_parallel_region_from_moe( output_total, expert_parallel=self.expert_parallel ) if self.add_bias: output_bias_total = \ reduce_scatter_to_sequence_parallel_region_from_moe( output_bias_total, expert_parallel=self.expert_parallel) # bias is duplicated across tensor parallelism ranks; # reduce scatter reduces bias across tensor parallel_ranks output_bias_total = \ output_bias_total/mpu.get_tensor_model_parallel_world_size() output_total = output_total*max_prob output_total = output_total.view(s, b, h) if self.add_bias: output_bias_total = output_bias_total*max_prob output_bias_total = output_bias_total.view(s, b, h) else: output_bias_total = None return output_total, output_bias_total class CoreAttention(MegatronModule): def __init__(self, layer_number, config, attn_mask_type=AttnMaskType.padding): super(CoreAttention, self).__init__() self.fp16 = config.fp16 self.bf16 = config.bf16 self.apply_query_key_layer_scaling = config.apply_query_key_layer_scaling self.attention_softmax_in_fp32 = config.attention_softmax_in_fp32 if self.apply_query_key_layer_scaling: self.attention_softmax_in_fp32 = True self.layer_number = max(1, layer_number) self.attn_mask_type = attn_mask_type self.sequence_parallel = config.sequence_parallel projection_size = config.kv_channels * config.num_attention_heads # Per attention head and per partition values. world_size = mpu.get_tensor_model_parallel_world_size() self.hidden_size_per_partition = core.utils.divide(projection_size, world_size) self.hidden_size_per_attention_head = core.utils.divide( projection_size, config.num_attention_heads) self.num_attention_heads_per_partition = core.utils.divide( config.num_attention_heads, world_size) coeff = None self.norm_factor = math.sqrt(self.hidden_size_per_attention_head) if self.apply_query_key_layer_scaling: coeff = self.layer_number self.norm_factor *= coeff self.scale_mask_softmax = FusedScaleMaskSoftmax( self.fp16, self.bf16, self.attn_mask_type, config.masked_softmax_fusion, attention_mask_func, self.attention_softmax_in_fp32, coeff) # Dropout. Note that for a single iteration, this layer will generate # different outputs on different number of parallel partitions but # on average it should not be partition dependent. self.attention_dropout = torch.nn.Dropout(config.attention_dropout) def forward(self, query_layer, key_layer, value_layer, attention_mask): # =================================== # Raw attention scores. [b, np, s, s] # =================================== # [b, np, sq, sk] output_size = (query_layer.size(1), query_layer.size(2), query_layer.size(0), key_layer.size(0)) # [sq, b, np, hn] -> [sq, b * np, hn] query_layer = query_layer.reshape(output_size[2], output_size[0] * output_size[1], -1) # [sk, b, np, hn] -> [sk, b * np, hn] key_layer = key_layer.view(output_size[3], output_size[0] * output_size[1], -1) # preallocting input tensor: [b * np, sq, sk] matmul_input_buffer = mpu.get_global_memory_buffer().get_tensor( (output_size[0]*output_size[1], output_size[2], output_size[3]), query_layer.dtype, "mpu") # Raw attention scores. [b * np, sq, sk] matmul_result = torch.baddbmm( matmul_input_buffer, query_layer.transpose(0, 1), # [b * np, sq, hn] key_layer.transpose(0, 1).transpose(1, 2), # [b * np, hn, sk] beta=0.0, alpha=(1.0/self.norm_factor)) # change view to [b, np, sq, sk] attention_scores = matmul_result.view(*output_size) - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - megatron_patch/model/mistral/transformer.py [39:391]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - try: from einops import rearrange except ImportError: rearrange = None try: from flash_attn.flash_attn_interface import flash_attn_unpadded_func except ImportError: try: from flash_attn.flash_attn_interface import flash_attn_varlen_func as flash_attn_unpadded_func except ImportError: flash_attn_unpadded_func = None """ We use the following notation throughout this file: h: hidden size n: number of attention heads p: number of model parallel partitions np: n/p hp: h/p hn: h/n b: batch size s: sequence length l: number of layers Transformer takes input of size [s, b, h] and returns a tensor of the same size. We use the following arguments: hyperparameters: transformer hyperparameters """ class DropPath(MegatronModule): """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). """ def __init__(self, drop_prob=0.): super(DropPath, self).__init__() self.drop_prob = drop_prob def forward(self, hidden_state): if self.drop_prob == 0. or not self.training: return hidden_state keep_prob = 1 - self.drop_prob # work with diff dim tensors, not just 2D ConvNets # hidden_state: [s, b, h] shape = (1,) + (hidden_state.shape[1],) + (1,) * (hidden_state.ndim - 2) random_tensor = keep_prob + \ torch.rand(shape, dtype=hidden_state.dtype, device=hidden_state.device) random_tensor.floor_() # binarize output = hidden_state.div(keep_prob) * random_tensor return output class ParallelMLP(MegatronModule): """MLP. MLP will take the input with h hidden state, project it to 4*h hidden dimension, perform nonlinear transformation, and project the state back into h hidden dimension. """ def __init__(self, config, is_expert=False): super(ParallelMLP, self).__init__() args = get_args() self.add_bias = config.add_bias_linear ffn_hidden_size = config.ffn_hidden_size if config.gated_linear_unit: ffn_hidden_size *= 2 # Project to 4h. If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf self.dense_h_to_4h = tensor_parallel.ColumnParallelLinear( config.hidden_size, ffn_hidden_size, config=config, init_method=config.init_method, bias=self.add_bias, gather_output=False, skip_bias_add=True, is_expert=is_expert, ) self.bias_gelu_fusion = False self.activation_func = None self.swiglu = args.swiglu if args.openai_gelu: self.activation_func = openai_gelu elif args.onnx_safe: self.activation_func = erf_gelu elif args.swiglu: def swiglu(x): x = torch.chunk(x, 2, dim=-1) return F.silu(x[0]) * x[1] self.activation_func = swiglu elif args.squared_relu: def squared_relu(x): return torch.pow(F.relu(x), 2) self.activation_func = squared_relu else: self.bias_gelu_fusion = args.bias_gelu_fusion self.activation_func = F.gelu # Project back to h. self.dense_4h_to_h = tensor_parallel.RowParallelLinear( config.ffn_hidden_size, config.hidden_size, config=config, init_method=config.output_layer_init_method, bias=self.add_bias, input_is_parallel=True, skip_bias_add=True, is_expert=is_expert, ) def forward(self, hidden_states): # [s, b, 4hp] intermediate_parallel, bias_parallel = self.dense_h_to_4h(hidden_states) if self.bias_gelu_fusion: assert self.add_bias is True assert self.activation_func == F.gelu intermediate_parallel = bias_gelu_impl(intermediate_parallel, bias_parallel) else: if bias_parallel is not None: intermediate_parallel = intermediate_parallel + bias_parallel intermediate_parallel = self.activation_func(intermediate_parallel) # [s, b, h] output, output_bias = self.dense_4h_to_h(intermediate_parallel) return output, output_bias def sinkhorn(cost, tol=0.0001): cost = torch.exp(cost) d0 = torch.ones(cost.size(0), device=cost.device, dtype=cost.dtype) d1 = torch.ones(cost.size(1), device=cost.device, dtype=cost.dtype) eps = 0.00000001 error = 1e9 d1_old = d1 while error > tol: d0 = (1/d0.size(0))*1/(torch.sum(d1*cost,1) + eps) d1 = (1/d1.size(0))*1/(torch.sum(d0.unsqueeze(1)*cost,0)+eps) error = torch.mean(torch.abs(d1_old-d1)) d1_old = d1 return d1*cost*d0.unsqueeze(1) class SwitchMLP(MegatronModule): """ Routes input to one of N MLP "experts" """ def __init__(self, config): super(SwitchMLP, self).__init__() args = get_args() self.router = torch.nn.Linear(args.hidden_size, args.num_experts) self.expert_parallel = config.expert_parallel self.sequence_parallel = config.sequence_parallel self.add_bias = config.add_bias_linear if self.expert_parallel: assert args.num_experts % mpu.get_data_parallel_world_size() == 0 self.num_local_experts = args.num_experts // mpu.get_data_parallel_world_size() local_expert_indices_offset = mpu.get_data_parallel_rank() * self.num_local_experts self.local_expert_indices = [local_expert_indices_offset + i for i in range(self.num_local_experts)] else: self.num_local_experts = args.num_experts self.local_expert_indices = [i for i in range(self.num_local_experts)] self.local_experts = torch.nn.ModuleList() for i in range(self.num_local_experts): self.local_experts.append(ParallelMLP(config, is_expert=True)) def gather_indices(self, local_indices): """ Gather tensors and concatinate along the first dimension.""" if self.expert_parallel: group = get_tensor_and_data_parallel_group() else: group = get_tensor_model_parallel_group() world_size = torch.distributed.get_world_size(group=group) # Bypass the function if we are using only 1 GPU. if world_size == 1: return local_indices dim_size = list(local_indices.size()) dim_size[0] = dim_size[0] * world_size # TODO pre allocate memory output = torch.empty(dim_size, dtype=local_indices.dtype, device=torch.cuda.current_device()) torch.distributed._all_gather_base( output, local_indices.contiguous(), group=group ) return output def forward(self, hidden_states): # hidden_states: [b, s, h] args = get_args() s = hidden_states.size(0) b = hidden_states.size(1) h = hidden_states.size(2) route = self.router(hidden_states).view(-1, args.num_experts) # TODO (rprenger) Right now we're just using the sinkhorn algorithm # for load balancing. There should be an option to do no load balancing # and the algorithm and parametets should be further tested if self.training: with torch.no_grad(): sinkroute = sinkhorn(route.detach().to(dtype=torch.float32)) _, max_ind = torch.max(sinkroute, dim=1) route = torch.sigmoid(route) max_prob = route[torch.arange(route.size(0)), max_ind] else: route = torch.sigmoid(route) max_prob, max_ind = torch.max(route, dim=1) max_prob = torch.unsqueeze(max_prob, 1) hidden_states = hidden_states.view(-1, hidden_states.size(2)) # TODO (rprenger) TODO this could be made easier to read # Converting [s, b, h] to [s*b, h]. # Each vector could be routed differently if self.sequence_parallel or self.expert_parallel: global_hidden_states = \ gather_from_sequence_parallel_region_to_moe( hidden_states, expert_parallel=self.expert_parallel ) global_indices = self.gather_indices(max_ind) else: global_hidden_states = hidden_states global_indices = max_ind output_total = torch.zeros_like(global_hidden_states) if self.add_bias: output_bias_total = torch.zeros_like(global_hidden_states) for expert_num, expert in enumerate(self.local_experts): local_expert_index = self.local_expert_indices[expert_num] local_indices = (global_indices == local_expert_index).nonzero() hidden = global_hidden_states[local_indices, :] output, output_bias = expert(hidden) output_total[local_indices, :] = output if self.add_bias: output_bias = output_bias.expand_as(output) output_bias_total[local_indices, :] = output_bias if self.sequence_parallel or self.expert_parallel: output_total = \ reduce_scatter_to_sequence_parallel_region_from_moe( output_total, expert_parallel=self.expert_parallel ) if self.add_bias: output_bias_total = \ reduce_scatter_to_sequence_parallel_region_from_moe( output_bias_total, expert_parallel=self.expert_parallel) # bias is duplicated across tensor parallelism ranks; # reduce scatter reduces bias across tensor parallel_ranks output_bias_total = \ output_bias_total/mpu.get_tensor_model_parallel_world_size() output_total = output_total*max_prob output_total = output_total.view(s, b, h) if self.add_bias: output_bias_total = output_bias_total*max_prob output_bias_total = output_bias_total.view(s, b, h) else: output_bias_total = None return output_total, output_bias_total class CoreAttention(MegatronModule): def __init__(self, layer_number, config, attn_mask_type=AttnMaskType.padding): super(CoreAttention, self).__init__() self.fp16 = config.fp16 self.bf16 = config.bf16 self.apply_query_key_layer_scaling = config.apply_query_key_layer_scaling self.attention_softmax_in_fp32 = config.attention_softmax_in_fp32 if self.apply_query_key_layer_scaling: self.attention_softmax_in_fp32 = True self.layer_number = max(1, layer_number) self.attn_mask_type = attn_mask_type self.sequence_parallel = config.sequence_parallel projection_size = config.kv_channels * config.num_attention_heads # Per attention head and per partition values. world_size = mpu.get_tensor_model_parallel_world_size() self.hidden_size_per_partition = core.utils.divide(projection_size, world_size) self.hidden_size_per_attention_head = core.utils.divide( projection_size, config.num_attention_heads) self.num_attention_heads_per_partition = core.utils.divide( config.num_attention_heads, world_size) coeff = None self.norm_factor = math.sqrt(self.hidden_size_per_attention_head) if self.apply_query_key_layer_scaling: coeff = self.layer_number self.norm_factor *= coeff self.scale_mask_softmax = FusedScaleMaskSoftmax( self.fp16, self.bf16, self.attn_mask_type, config.masked_softmax_fusion, attention_mask_func, self.attention_softmax_in_fp32, coeff) # Dropout. Note that for a single iteration, this layer will generate # different outputs on different number of parallel partitions but # on average it should not be partition dependent. self.attention_dropout = torch.nn.Dropout(config.attention_dropout) def forward(self, query_layer, key_layer, value_layer, attention_mask): # =================================== # Raw attention scores. [b, np, s, s] # =================================== # [b, np, sq, sk] output_size = (query_layer.size(1), query_layer.size(2), query_layer.size(0), key_layer.size(0)) # [sq, b, np, hn] -> [sq, b * np, hn] query_layer = query_layer.reshape(output_size[2], output_size[0] * output_size[1], -1) # [sk, b, np, hn] -> [sk, b * np, hn] key_layer = key_layer.view(output_size[3], output_size[0] * output_size[1], -1) # preallocting input tensor: [b * np, sq, sk] matmul_input_buffer = mpu.get_global_memory_buffer().get_tensor( (output_size[0]*output_size[1], output_size[2], output_size[3]), query_layer.dtype, "mpu") # Raw attention scores. [b * np, sq, sk] matmul_result = torch.baddbmm( matmul_input_buffer, query_layer.transpose(0, 1), # [b * np, sq, hn] key_layer.transpose(0, 1).transpose(1, 2), # [b * np, hn, sk] beta=0.0, alpha=(1.0/self.norm_factor)) # change view to [b, np, sq, sk] attention_scores = matmul_result.view(*output_size) - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -