Unable to Reproduce Llamba Paper's Figure 4 Throughput Results at Large Batch Sizes

[shambhavicodes]

Issue Description

I'm trying to reproduce Figure 4 from the Llamba paper which shows throughput scaling with batch sizes up to 2048, but I'm encountering persistent Out of Memory (OOM) errors even on H100 GPU when running with moderate batch sizes.

Environment

• GPU: NVIDIA H100 (80GB)
• Model: Llamba-8B
• Sequence length tested: 8192 and 4096
• Using the provided benchmarking script

Observations

• Float32 Precision: OOM errors at batch size 16 on H100
• BFloat16 Precision: OOM errors at batch size 32 on H100
• Issue persists regardless of sequence length (tested with both 8192 and 4096)

This severely limits my ability to reproduce the throughput scaling shown in Figure 4 of the paper, where batch sizes up to 2048 are reported.

Code Used for Benchmarking

import time
import os
import torch
from transformers import AutoTokenizer
from cartesia_pytorch.Llamba.llamba import LlambaLMHeadModel

device = "cuda" if torch.cuda.is_available() else "cpu"
os.environ["TOKENIZERS_PARALLELISM"] = "false"

BATCH_SIZES = [32]  # Tried with [1, 8, 16, 32]
PROMPT_LEN = 1
GEN_LEN = 8192

@torch.inference_mode()
def warmup(generate_fn):
    for _ in range(3):
        _ = generate_fn()
    torch.cuda.synchronize()

@torch.inference_mode()
def time_bench(generate_fn, batch_size):
    torch.cuda.synchronize()
    start = time.time()
    out = generate_fn()
    torch.cuda.synchronize()
    elapsed = time.time() - start

    mem_allocated = torch.cuda.memory_allocated() / (1024**3)
    max_mem_allocated = torch.cuda.max_memory_allocated() / (1024**3)

    print(f"\n[Batch size: {batch_size}] | Prompt len: {PROMPT_LEN} | Gen len: {GEN_LEN}")
    print(f"Throughput: {(batch_size * GEN_LEN) / elapsed:.1f} tokens/sec")
    print(f"Time per run: {elapsed * 1000:.1f} ms")
    print(f"CUDA memory allocated: {mem_allocated:.2f} GB")
    print(f"CUDA max memory allocated: {max_mem_allocated:.2f} GB")

@torch.inference_mode()
def main():
    model = LlambaLMHeadModel.from_pretrained("cartesia-ai/Llamba-8B")
    tokenizer = AutoTokenizer.from_pretrained("/mnt/checkpoints/Llama-3.1-8B")    
    tokenizer.pad_token = tokenizer.eos_token

    model.to(device=device, dtype=torch.bfloat16)
    model = torch.compile(model, fullgraph=True)
    model.eval()

    for bs in BATCH_SIZES:
        input_ids = torch.randint(0, tokenizer.vocab_size, (bs, PROMPT_LEN), dtype=torch.long, device=device)

        # CUDA Graph Capture
        static_input_ids = input_ids.clone()
        static_max_len = static_input_ids.shape[1] + GEN_LEN

        def _gen_step():
            return model.generate(
                input_ids=static_input_ids,
                max_length=static_max_len,
                cg=True,
                return_dict_in_generate=False,
                output_scores=False,
                eos_token_id=tokenizer.eos_token_id,
            )

        warmup(_gen_step)
        time_bench(_gen_step, bs)
        torch.cuda.reset_peak_memory_stats()

if __name__ == "__main__":
    main()

Question

Could you please provide guidance on how to reproduce the throughput measurements shown in Figure 4 of the Llamba paper, which demonstrates scaling to batch sizes up to 2048? What modifications to the benchmarking approach or hardware setup are needed to achieve such scaling?

[dandenkijin]

Would no response indicate that this project is dead, or . . .
something else?