Get the memory (in bytes) required to store the embedding layer, given the number of parameters in the embedding layer, the data type (defaults to FP32) used for the weights, and the tensor parallelism size (Megatron-LM partitions the embedding layer across the tensor paralle
(
self,
ds_zero: DSZeRO = DSZeRO.NONE,
)
| 516 | return memory_optimizer_state_last_layernorm, memory_gradient_last_layernorm |
| 517 | |
| 518 | def get_memory_embedding( |
| 519 | self, |
| 520 | ds_zero: DSZeRO = DSZeRO.NONE, |
| 521 | ) -> float: |
| 522 | """Get the memory (in bytes) required to store the embedding layer, given the |
| 523 | number of parameters in the embedding layer, the data type (defaults to FP32) |
| 524 | used for the weights, and the tensor parallelism size (Megatron-LM partitions |
| 525 | the embedding layer across the tensor parallel groups). |
| 526 | |
| 527 | Args: |
| 528 | ds_zero (DSZeRO, optional): which DeepSpeed ZeRO stage to use. Defaults to DSZeRO.NONE (disabled, no sharding). |
| 529 | |
| 530 | Returns: |
| 531 | float: the memory (in bytes) required to store the embedding layer |
| 532 | """ |
| 533 | dtype_bytes = self.dtype_config.embedding_bits / BITS_PER_BYTE |
| 534 | memory_embedding = (self.get_num_params_embedding() / |
| 535 | self.parallelism_config.tp_size) * dtype_bytes |
| 536 | if ds_zero == DSZeRO.STAGE_3: |
| 537 | memory_embedding /= self.parallelism_config.dp_size |
| 538 | |
| 539 | return memory_embedding |
| 540 | |
| 541 | def get_activation_memory_per_layer_attn( |
| 542 | self, |
no test coverage detected