DOCSEARCHbyAlgolia
Submit Docs

LLMCompass

repository·main·Indexed 42 days ago

https://github.com/princetonuniversity/llmcompass

An implementation of hardware design enablement for Large Language Model (LLM) inference, as presented in the ISCA 2024 paper. It provides tools for architecture exploration (AE) and simulation, allowing users to define hardware configurations via JSON and build custom LLM computational graphs using primitive operators like matmul, softmax, and allreduce.

Tokens
2.6K
Snippets
5
Records
8
Agent score
5%

What's inside LLMCompass

  1. Build a custom LLM computational graph

    main

    You can build custom LLM computational graphs by inheriting from the Operator class. This allows you to define custom logic using the provided primitive operators:

    • matmul
    • softmax
    • layernorm
    • gelu
    • allreduce

    When implementing a custom class, you must define:

    • __init__: Initialize the required operators.
    • __call__: Build the computational graph. Tensor shapes are automatically calculated for the simulation.
    • compile_and_simulate: Simulates the operators to return total latency and other runtimes.
    • roofline_model (optional): Provides a roofline model analysis.
    • run_on_gpu (optional): Executes the graph on real GPUs using PyTorch.

  4. Run AE Experiment workflows

    main

    The AE (Architecture Exploration) experiments are organized by figure number. Navigate to the specific figure directory within ae/ and execute the corresponding .sh script. Results (figures) are generated within the directory where the script is run. You can compare your output with the reference results located in ae\expected_results.

    # Figure 5 (approx. 100 min)
$ cd ae/figure5
$ bash run_figure5.sh 

# Figure 6 (approx. 1 min)
$ cd ae/figure6
$ bash run_figure6.sh

# Figure 7 (approx. 20 min)
$ cd ae/figure7
$ bash run_figure7.sh

# Figure 8 (approx. 40 min)
$ cd ae/figure8
$ bash run_figure8.sh

# Figure 9 (approx. 30 min)
$ cd ae/figure9
$ bash run_figure9.sh

# Figure 10 (approx. 45 min)
$ cd ae/figure10
$ bash run_figure10.sh

# Figure 11 (approx. 5 min)
$ cd ae/figure11
$ bash run_figure11.sh

# Figure 12 (approx. 4 hours)
$ cd ae/figure12
$ bash run_figure12.sh

  5. Run a LLMCompass simulation

    main

    To run a full simulation, follow these three steps:

    1. Load Hardware: Use read_architecture_template and template_to_system to parse your JSON configuration.
    2. Instantiate Model: Create your model (e.g., TransformerBlockAutoRegressionTP) with parameters like d_model, n_heads, device_count (for Tensor Parallelism), and data_type.
    3. Simulate: Call compile_and_simulate on the model instance, passing the system object and a simulation strategy (e.g., "heuristic-GPU").
    from design_space_exploration.dse import template_to_system, read_architecture_template

# 1. Load Hardware
specs = read_architecture_template("PATH/TO/YOUR/JSON")
system = template_to_system(specs)

# 2. Instantiate Model
model_auto_regression = TransformerBlockAutoRegressionTP(
        d_model=12288,
        n_heads=96,
        device_count=1,
        data_type=data_type_dict["fp16"],
    )
_ = model_auto_regression(
	Tensor([bs, 1, 12288], data_type_dict["fp16"]),
	seq_len,
)

# 3. Run Simulation
auto_regression_latency_simulated = model_auto_regression.compile_and_simulate(
	system, "heuristic-GPU"
)

  6. Set up the LLMCompass environment with Conda

    main

    To prepare the environment for LLMCompass, create a new Conda environment with Python 3.9 and install the required dependencies including scalesim, pytorch, matplotlib, seaborn, and scipy.

    $ conda create -n llmcompass_ae python=3.9
$ conda activate llmcompass_ae
$ pip3 install scalesim
$ conda install pytorch==2.0.0 -c pytorch
$ pip3 install matplotlib
$ pip3 install seaborn
$ pip3 install scipy

  7. Install LLMCompass via GitHub

    main

    Clone the ISCA_AE branch from the official repository and initialize the submodules to ensure all dependencies are present.

    $ git clone -b ISCA_AE https://github.com/PrincetonUniversity/LLMCompass
$ cd LLMCompass
$ git submodule init
$ git submodule update --recursive

  8. Configure a Hardware Configuration JSON

    main

    To run a simulation, you must first define a hardware configuration in a JSON file. This file describes the node's device count, interconnect topology, and detailed device specifications including compute chiplets, memory protocols, and I/O capabilities.

    Key configuration knobs include:

    • interconnect: Defines the link properties (bandwidth, latency, flit size) and topology (currently supports FC for fully-connected and RING).
    • device.compute_chiplet: Contains core_count (for performance modeling) and physical_core_count (for area modeling).
    • device.io: Specifies global_buffer_bandwidth_per_cycle_byte and memory channel counts.
    • device.memory: Defines total_capacity_GB and memory_protocol (supported: HBM2e, DDR4, DDR5, PCIe4, PCIe5).
    {
    "name": "NVIDIA A100(80GB)x4", 
    "device_count": 4,
    "interconnect": {
        "link": {
            "name": "NVLink3",
            "bandwidth_per_direction_byte": 25e9,
            "bandwidth_both_directions_byte": 50e9,
            "latency_second": 8.92e-6,
            "flit_size_byte": 16,
            "header_size_byte": 16,
            "max_payload_size_byte": 256
        },
        "link_count_per_device": 12,
        "topology": "FC"
    },
    "device": {
        "frequency_Hz": 1410e6,
        "compute_chiplet_count": 1,
        "compute_chiplet": {
            "physical_core_count": 128,
            "core_count": 128,
            "process_node": "7nm",
            "core": {
                "sublane_count": 4,
                "systolic_array": {
                    "array_width": 16,
                    "array_height": 16,
                    "data_type": "fp16",
                    "mac_per_cycle": 1
                },
                "vector_unit": {
                    "vector_width": 32,
                    "flop_per_cycle": 4,
                    "data_type": "fp16",
                    "int32_count": 16,
                    "fp16_count": 0,
                    "fp32_count": 16,
                    "fp64_count": 8
                },
                "register_file": {
                    "num_reg_files": 1,
                    "num_registers": 16384,
                    "register_bitwidth":32,
                    "num_rdwr_ports":4
                },
                "SRAM_KB": 192
            }
        },
        "memory_protocol": "HBM2e",
        "io": {
            "process_node": "7nm",
            "global_buffer_MB": 48,
            "physical_global_buffer_MB": 48,
            "global_buffer_bandwidth_per_cycle_byte": 5120,
            "memory_channel_physical_count": 6,
            "memory_channel_active_count": 5,
            "pin_count_per_channel": 1024,
            "bandwidth_per_pin_bit": 3.2e9
        },
        "memory": {
            "total_capacity_GB": 80
        }
    }
}