Get the latency for the forward pass of a transformer layer, given the batch size, sequence length, training or inference, activation recomputation strategy, and layernorm data type. The latency is the sum of the latency for the attention module, MLP module, two layernorms, a
(
self,
batch_size: int,
seq_len: int,
is_inference: bool = True,
activation_recomputation:
ActivationRecomputation = ActivationRecomputation.NONE,
layernorm_dtype_bytes: int = BYTES_FP32,
)
| 1164 | ) |
| 1165 | |
| 1166 | def get_latency_fwd_per_layer( |
| 1167 | self, |
| 1168 | batch_size: int, |
| 1169 | seq_len: int, |
| 1170 | is_inference: bool = True, |
| 1171 | activation_recomputation: |
| 1172 | ActivationRecomputation = ActivationRecomputation.NONE, |
| 1173 | layernorm_dtype_bytes: int = BYTES_FP32, |
| 1174 | ) -> tuple: |
| 1175 | """Get the latency for the forward pass of a transformer layer, given the batch |
| 1176 | size, sequence length, training or inference, activation recomputation strategy, |
| 1177 | and layernorm data type. The latency is the sum of the latency for the attention |
| 1178 | module, MLP module, two layernorms, and two (Megatron-LM tp implementation) |
| 1179 | allreduce communications across the tensor parallel group. |
| 1180 | |
| 1181 | Args: |
| 1182 | batch_size (int): batch size |
| 1183 | seq_len (int): sequence length |
| 1184 | is_inference (bool, optional): whether it is inference or not. Defaults to True. |
| 1185 | activation_recomputation (ActivationRecomputation, optional): activation recomputation strategy. Defaults to ActivationRecomputation.NONE. |
| 1186 | layernorm_dtype_bytes (int, optional): number of bytes in the data type for the layernorm activations. Defaults to BYTES_FP32. Often has to be FP32 in training to maintain model accuracy. |
| 1187 | |
| 1188 | Returns: |
| 1189 | tuple: a tuple of the latency in seconds for the forward pass of a transformer layer and its breakdown dict |
| 1190 | """ |
| 1191 | latency_fwd_per_layer_attn = self.get_latency_fwd_per_layer_attn( |
| 1192 | batch_size, seq_len, is_inference, activation_recomputation) |
| 1193 | |
| 1194 | latency_fwd_per_layer_mlp = self.get_latency_fwd_per_layer_mlp( |
| 1195 | batch_size, seq_len, is_inference, activation_recomputation) |
| 1196 | |
| 1197 | latency_fwd_per_layer_layernorm = ( |
| 1198 | self.get_latency_fwd_per_layer_layernorm( |
| 1199 | batch_size, |
| 1200 | seq_len, |
| 1201 | activation_recomputation, |
| 1202 | layernorm_dtype_bytes, |
| 1203 | )) |
| 1204 | |
| 1205 | latency_fwd_per_layer_tp_comm = self.get_latency_fwd_per_layer_tp_comm( |
| 1206 | batch_size, |
| 1207 | seq_len, |
| 1208 | self.dtype_config.activation_bits / BITS_PER_BYTE, |
| 1209 | ) |
| 1210 | |
| 1211 | latency_per_layer = (latency_fwd_per_layer_attn + |
| 1212 | latency_fwd_per_layer_mlp + |
| 1213 | 2 * latency_fwd_per_layer_layernorm + |
| 1214 | 2 * latency_fwd_per_layer_tp_comm) |
| 1215 | |
| 1216 | logger.debug("latency_fwd_per_layer_layernorm:" |
| 1217 | f" {round(latency_fwd_per_layer_layernorm*1000, 3)} ms," |
| 1218 | " latency_fwd_per_layer_tp_comm:" |
| 1219 | f" {round(latency_fwd_per_layer_tp_comm*1000, 3)} ms") |
| 1220 | |
| 1221 | logger.debug( |
| 1222 | f"latency_per_layer: {round(latency_per_layer*1000, 3)} ms" |
| 1223 | f" ({round(latency_fwd_per_layer_attn*1000, 3)} +" |
no test coverage detected