Function _initialize_mem_buffs

codegeex/megatron/initialize.py:331–337

Initialize manually allocated static memory.

()

def _initialize_mem_buffs():
    """Initialize manually allocated static memory."""
    args = get_args()

    # Initialize memory for checkpointed activations.
    if args.distribute_checkpointed_activations:
        mpu.init_checkpointed_activations_memory_buffer()