🚀 The feature, motivation and pitch

motivation

The pin memory allocation module adopts a cold-up strategy, with an initially empty cache. During the first few training iterations of the model, frequent requests for memory resource allocation are often required, leading to prolonged training time and reduced training efficiency.

solution

To address this issue, we propose a pin memory management mechanism based on warm-up and buddy system to resolve the cold-start performance problem during the early stages of model training.

warm-up

Pre-allocate multiple memory blocks, with the quantity and size controllable via environment variables. By default, 8 memory blocks of 1 GB each are pre-allocated.

buddy system

When a memory request is made, if the free list for the corresponding size is empty, the system will first search in the free list for larger sizes. If a suitable block is found, it will be split, and the resulting smaller block will either be allocated or inserted into the corresponding free list. When a memory block is released, the system will search the free list for its buddy block and merge them to satisfy subsequent memory requests.

threshold control

The buddy system's splitting mechanism inevitably causes memory fragmentation, primarily from splitting triggered by small-block requests. Experimental results show that memory request time is roughly proportional to block size, meaning KB-level requests account for a negligible time overhead (<1%) during training. Thus, we introduce a threshold control mechanism, configurable via environment variables, that sets a minimum block size for buddy system allocation. Requests smaller than this threshold are handled by the original allocation mechanism, substantially reducing the fragmentation rate.

adaptive clear mechanism

When calling empty_cache to release memory, only complete pre-allocated blocks are returned; split blocks with unreleased portions are not returned. When a task that previously called empty_cache is run again, the previously returned memory blocks are reallocated to the task.

usage

set pre-allocated block size (unit: MB, default 1024)

export BUDDY_BLOCK_SIZE_MB=512

set pre-allocated block count (default 8)

export BUDDY_BLOCK_COUNT=16

set the threshold for memory requests allocated via the buddy system (unit: MB, default 4)

export BUDDY_THRESHOLD_MB=2

experiment

We conducted comparative experiments training the DeepSeek-V3 model on 4 machines with 32 H100 GPUs. Experimental results show that this solution reduces memory allocation time by over 40% and total training time by 11%. Additionally, the threshold mechanism effectively prevents memory fragmentation.

S: pre-allocated block size C: pre-allocated block count T: threshold

Currently, this optimization solution has been officially deployed internally at Kuaishou, demonstrating strong practicality.

Alternatives

No response

Additional context

No response

cc @jerryzh168 @ptrblck @msaroufim @eqy @tinglvv @nWEIdia