Reduce Memory Usage During Inference

[bigximik]

🎯 Goal (What & Why)

The current memory footprint during the inference forward pass is approximately 2.5× higher for the same model and batch size when using Fast-LLM compared to Hugging Face Transformers.

🧠 Qwen2 1.5B — Batch Size: 16×4096, Flash Attention, bfloat16, H100 GPU

Test	Peak GPU Memory Usage (MB)
HF (no loss calculation)	22,162.28
HF (with loss calculation)	40,962.78
Fast-LLM (no loss calculation)	59,013.70
Fast-LLM (with loss calculation)	OOM

What is a reasonable target for reducing Fast-LLM's memory usage?

🚀 Execution Plan

(This section may start as an incomplete draft but must be defined before implementation begins.)

Step 1: What is the smallest working version?

(Describe the simplest way to implement this feature with minimal effort.)

Step 2: What additional optimizations are possible (but optional)?

(List potential refinements that can be added in later PRs if needed.)

📌 Acceptance Criteria (Must-Haves for Completion)

• The feature must be functional and tested.
• The implementation must be documented in practical terms.
• The PR must include a performance/impact summary.
• No refactors unless directly necessary for feature completion.

🛠️ Project Management

• Assign the project to the Fast-LLM project.
• Set the Estimate field (in days) in the GitHub project.
• Use the Size field to categorize the PR size (Small/Medium/Large).
• Assign an owner when opening the issue.

[bigximik]

This is the code used for testing:

import transformers
import torch
import torch.nn.functional as F


from fast_llm.models.gpt.config import GPTModelConfig, GPTBatchConfig
from fast_llm.engine.schedule.config import ScheduleConfig
from fast_llm.engine.schedule.runner import ScheduleRunner
from fast_llm.engine.distributed.distributed import Distributed
from fast_llm.engine.schedule.schedule import Schedule
from fast_llm.engine.distributed.config import PhaseType
from fast_llm.config import NoAutoValidate
from fast_llm.engine.checkpoint.config import CheckpointLoadConfig
from fast_llm.data.data.gpt.data import GPTBatch


def run_hf(model_path: str, do_loss: bool):
    tokenizer = transformers.AutoTokenizer.from_pretrained(model_path)
    model = transformers.AutoModelForCausalLM.from_pretrained(
        model_path,
        torch_dtype=torch.bfloat16,  # or torch.bfloat16 if you prefer
        attn_implementation="flash_attention_2",  # Requires supported model & transformers>=4.35
    ).cuda()

    seq_len = 4097
    batch_size = 16
    data = torch.randint(0, tokenizer.vocab_size, (batch_size, seq_len), device="cuda")

    with torch.no_grad():
        input_ids = data[:, :-1].contiguous()  # [16, 4096]
        if do_loss:
            labels = data[:, 1:].contiguous()
        res = model(input_ids=input_ids, use_cache=False, return_dict=True)
        logits = res.logits

        if do_loss:
            # Flatten for loss computation
            logits = logits.view(-1, logits.size(-1))  # [batch_size * (seq_len-1), vocab_size]
            labels = labels.view(-1)  # [batch_size * (seq_len-1)]

            # Compute loss
            loss = F.cross_entropy(logits, labels, ignore_index=tokenizer.pad_token_id)
            print("Cross-entropy loss:", loss.item())


def run_fast_llm(model_path: str, do_loss: bool):
    tokenizer = transformers.AutoTokenizer.from_pretrained(model_path)
    updates = {}
    updates[("base_model", "transformer", "use_flash_attention")] = True
    updates[("base_model", "cross_entropy_impl")] = "fused"
    updates[("distributed", "training_dtype")] = "bf16"

    config = GPTModelConfig.from_dict({}, updates)

    with NoAutoValidate():
        checkpoint_load_config = CheckpointLoadConfig(path=model_path, model_weights=True, format="qwen2")
        checkpoint_load_config.setup(config)
        checkpoint_load_config.validate()

    config = config.from_pretrained(checkpoint_load_config, config)

    multi_stage = config.get_model_class()(config)
    schedule_config = ScheduleConfig()
    runner = ScheduleRunner(
        config=schedule_config,
        multi_stage=multi_stage,
        distributed_config=config.distributed,
    )
    distributed = Distributed(config.distributed)

    with torch.no_grad():
        multi_stage.setup(distributed)
        runner.setup(distributed)

    multi_stage.load_checkpoint(checkpoint_load_config)

    with NoAutoValidate():
        batch_config = GPTBatchConfig(micro_batch_size=16, sequence_length=4096, batch_size=16)
        batch_config.setup(config.distributed)
        batch_config.validate()

    schedule = Schedule(
        multi_stage=multi_stage,
        batch_config=batch_config,
        schedule_config=runner.config,
        distributed_config=distributed.config,
        phase=PhaseType.validation if do_loss else PhaseType.inference,
    )

    seq_len = 4097 if do_loss else 4096
    batch_size = 16
    data = torch.randint(0, tokenizer.vocab_size, (batch_size, seq_len), device="cuda")

    with torch.no_grad():
        input_ids = data
        res = runner.run_step(iter((GPTBatch(token_ids = input_ids),)), schedule, iteration=0)


if __name__ == "__main__":
    run_hf("/mnt/checkpoints/pretrained_models/Qwen2-1.5B-Instruct", True)
    #run_fast_llm("/mnt/checkpoints/pretrained_models/Qwen2-1.5B-Instruct", False)
    peak_mem = torch.cuda.max_memory_allocated()
    print(f"Peak GPU memory usage: {peak_mem / (1024**2):.2f} MB")

[bigximik]

This memory optimization would benefit standalone evaluations, since during training we will most likely use the same batch size per GPU for evaluation as we do for training—so the memory footprint will remain unchanged.