Get the latency for the forward pass of the transformer, given the batch size, sequence length, and whether it is inference or not, the activation recomputation strategy, and the number of bytes in the data type for the layernorm activations. Args: batch_
(
self,
batch_size: int,
seq_len: int,
is_inference: bool = True,
activation_recomputation:
ActivationRecomputation = ActivationRecomputation.NONE,
layernorm_dtype_bytes: int = BYTES_FP32,
breakdown_prefix: str = "",
)
| 1276 | return compute_latency |
| 1277 | |
| 1278 | def get_latency_fwd( |
| 1279 | self, |
| 1280 | batch_size: int, |
| 1281 | seq_len: int, |
| 1282 | is_inference: bool = True, |
| 1283 | activation_recomputation: |
| 1284 | ActivationRecomputation = ActivationRecomputation.NONE, |
| 1285 | layernorm_dtype_bytes: int = BYTES_FP32, |
| 1286 | breakdown_prefix: str = "", |
| 1287 | ) -> tuple: |
| 1288 | """Get the latency for the forward pass of the transformer, given the batch |
| 1289 | size, sequence length, and whether it is inference or not, the activation |
| 1290 | recomputation strategy, and the number of bytes in the data type for the |
| 1291 | layernorm activations. |
| 1292 | |
| 1293 | Args: |
| 1294 | batch_size (int): batch size |
| 1295 | seq_len (int): sequence length |
| 1296 | is_inference (bool, optional): whether it is inference or not. Defaults to True. |
| 1297 | activation_recomputation (ActivationRecomputation, optional): activation recomputation strategy. Defaults to ActivationRecomputation.NONE. |
| 1298 | 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. |
| 1299 | breakdown_prefix (str, optional): prefix for the breakdown dict keys. Defaults to "". |
| 1300 | Returns: |
| 1301 | tuple: a tuple of the latency in seconds for the forward pass of the transformer and its breakdown dict |
| 1302 | """ |
| 1303 | num_layers_per_gpu = int(self.model_config.num_layers / |
| 1304 | self.parallelism_config.pp_size) |
| 1305 | |
| 1306 | ( |
| 1307 | latency_fwd_per_layer, |
| 1308 | breakdown_per_layer, |
| 1309 | ) = self.get_latency_fwd_per_layer( |
| 1310 | batch_size, |
| 1311 | seq_len, |
| 1312 | is_inference, |
| 1313 | activation_recomputation, |
| 1314 | layernorm_dtype_bytes, |
| 1315 | ) |
| 1316 | |
| 1317 | latency_fwd_all_layers = latency_fwd_per_layer * num_layers_per_gpu |
| 1318 | |
| 1319 | latency_fwd_input_embedding = self.get_latency_fwd_input_embedding( |
| 1320 | batch_size, |
| 1321 | seq_len, |
| 1322 | dtype_bytes=self.dtype_config.embedding_bits / BITS_PER_BYTE, |
| 1323 | ) |
| 1324 | |
| 1325 | latency_fwd_output_embedding_loss = ( |
| 1326 | self.get_latency_fwd_output_embedding_loss(batch_size, seq_len)) |
| 1327 | |
| 1328 | total_latency = (latency_fwd_all_layers + latency_fwd_input_embedding + |
| 1329 | latency_fwd_output_embedding_loss) |
| 1330 | |
| 1331 | logger.debug("latency_fwd_all_layers:" |
| 1332 | f" {round(latency_fwd_all_layers*1000, 3)} ms" |
| 1333 | f" ({round(latency_fwd_per_layer*1000, 3)} ms x" |
| 1334 | f" {num_layers_per_gpu}), latency_fwd_input_embedding:" |
| 1335 | f" {round(latency_fwd_input_embedding*1000, 3)} ms," |
no test coverage detected