Llama3-8B (Text Generation)¶

Installation¶

1
2

$ cd rbln-model-zoo/ATOM/huggingface/transformers/text2text-generation/llama/llama3-8b
$ pip3 install -r requirements.txt

Compile Model and Activate RBLN Profiler¶

1
2
3

# Sample Text: "Hey, are you conscious? Can you talk to me?”
$ python3 compile.py
$ RBLN_PROFILER=1 python3 inference.py

Or, set rbln_activate_profiler=True to enable the RBLN Profiler:

1
2
3
4
5
6
7
8
9

# rbln-model-zoo/ATOM/huggingface/transformers/text2text-generation/llama/llama3-8b/inference.py
...
# Load compiled model
model = RBLNLlamaForCausalLM.from_pretrained(
    model_id=os.path.basename(model_id),
    export=False,
    rbln_activate_profiler=True,
)
...

Extract Profiled Data with Inference¶

Run inference to obtain profiling data:

1
2
3
4
5
6
7
8
9

# rbln-model-zoo/ATOM/huggingface/transformers/text2text-generation/llama/llama3-8b/inference.py
...
# Generate tokens
output_sequence = model.generate(
    inputs.input_ids,
    attention_mask=inputs.attention_mask,
    max_length=8192,
)
...

Alternatively, use scoped profiling methods as shown below:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

# rbln-model-zoo/ATOM/huggingface/transformers/text2text-generation/llama/llama3-8b/inference.py
...
from rebel.profiler import profile

with profile(output_dir="./profile_out"):
    # Generate tokens
    output_sequence = model.generate(
        inputs.input_ids,
        attention_mask=inputs.attention_mask,
        max_length=8192,
    )
...

LlamaDecoderLayer¶

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

# https://github.com/huggingface/transformers/blob/5fa35344755d8d9c29610b57d175efd03776ae9e/src/transformers/models/llama/modeling_llama.py#L332
class LlamaDecoderLayer(nn.Module):
    ...
    def forward(...) -> ...:
        ...

        hidden_states = self.input_layernorm(hidden_states)

        # Self Attention
        hidden_states, self_attn_weights = self.self_attn(
            hidden_states=hidden_states,
            ...
        )
        ...

LlamaAttention¶

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

# https://github.com/huggingface/transformers/blob/5fa35344755d8d9c29610b57d175efd03776ae9e/src/transformers/models/llama/modeling_llama.py#L269
class LlamaAttention(nn.Module):
    ...
    def forward(
        self,
        hidden_states: torch.Tensor,
        ...
    ) -> ...:
        ...
        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
        ...
        attn_output = self.o_proj(attn_output)
        ...

Layernorm Execution in Root Node¶

The 0_causal_lm.model.0.input_layernorm layer is small enough to run on a single device. Thus, this layer is processed on a single device and the output is transferred to other devices. While the root ATOM™ device processes input_layernorm, other devices use Neural DMA to load weight parameters that are needed later. After input_layernorm is computed, the output tensor is transferred to leaf ATOM™ devices with Device HDMA commands.

Send Output Data and Synchronization¶

The output tensor of input_layernorm is transferred to the DRAM of ATOM™ 0 by Task DMA. Also, because leaf ATOM™ devices should use the same tensor, it is transferred to them by Device HDMA. These processes occur with almost no time delay, allowing the leaf nodes to quickly process the next layer. When tensor data is transferred by a Device HDMA command, a receiver (leaf) ATOM™ executes the Device Sync command for integrity. After this command is processed, each device loads the data in the same way as data transferred from External HDMA. Note that the Device HDMA command is executed on the sender, not the receiver.

Gather data to root node¶

The Device HDMA command is executed not only when an output tensor computed in the root ATOM™ device is scattered, but also when output tensors computed in leaf ATOM™ devices are gathered. The above example shows how the output tensor computed on a leaf ATOM™ device is transferred to the root ATOM™ device. Each output tensor that is transferred by Device HDMA is transferred to the internal Shared Memory of the root ATOM™ by Task DMA. The Neural Engine Clusters command that gathers the output tensors depends on the final Task DMA command, which retrieves the last output tensor from a leaf ATOM™ device, so it runs only after that Task DMA command completes.

Naming policy about attention¶

In transformer-based networks, it is important to understand how the attention mechanism works. However, because existing frameworks mainly construct models using modules, it is hard to identify when and how attention operates. To help users understand, the RBLN Profiler appends an attention suffix to commands that consist of an attention operation.