code: def get_act_dict(
      self,
      tokenizer,
      dataset,
      num_samples=512,
      seq_len=512,
  ):
      llama_model = self.model
  
      llama_model.eval()
      device = next(llama_model.parameters()).device
      # print("model:", llama_model)
      act_dict = defaultdict(dict)
  
      def stat_io_hook(m, x, y, name):
          if isinstance(x, tuple):
              x = x[0]
          if name not in act_dict or "input" not in act_dict[name]:
              act_dict[name]["input"] = x.detach().abs().max().item()
          else:
              act_dict[name]["input"] = max(
                  act_dict[name]["input"], x.detach().abs().max().item()
              )
          if isinstance(y, tuple):
              y = y[0]
          if name not in act_dict or "output" not in act_dict[name]:
              act_dict[name]["output"] = y.detach().abs().max().item()
          else:
              act_dict[name]["output"] = max(
                  act_dict[name]["output"],
                  y.detach().abs().max().item(),
              )
  
      for name, m in llama_model.named_modules():
          if isinstance(m, LlamaAttention):
              setattr(m, "q_apply_rotary", LlamaApplyRotary())
              setattr(m, "k_apply_rotary", LlamaApplyRotary())
              m.forward = types.MethodType(
                  llama_decoder_layer_forward, m
              )
  
      hooks = []
      for name, m in llama_model.named_modules():
          if isinstance(m, LlamaApplyRotary):
              hooks.append(
                  m.register_forward_hook(
                      partial(stat_io_hook, name=name)
                  )
              )
          if isinstance(m, torch.nn.Linear):
              hooks.append(
                  m.register_forward_hook(
                      partial(stat_io_hook, name=name)
                  )
              )
  
      self.collect_act_dict(
          llama_model,
          tokenizer,
          dataset,
          act_dict,
          device,
          num_samples,
          seq_len,
      )
  
      for hook in hooks:
          hook.remove()
      return act_dict
  
imports: import math
  import os
  import types
  from collections import defaultdict
  from functools import partial
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.nn
  import torch.nn.functional
  from transformers import PreTrainedModel
  from transformers.modeling_outputs import BaseModelOutputWithPast
  from transformers.models.llama.configuration_llama import LlamaConfig
  from transformers.models.llama.modeling_llama import LLAMA_INPUTS_DOCSTRING, LlamaAttention, LlamaDecoderLayer, LlamaMLP, LlamaRotaryEmbedding, rotate_half
  from transformers.utils import add_start_docstrings_to_model_forward
  from colossalai.inference.kv_cache.batch_infer_state import BatchInferState
  from colossalai.kernel.triton import copy_kv_cache_to_dest, int8_rotary_embedding_fwd, smooth_llama_context_attn_fwd, smooth_token_attention_fwd
  from .base_model import BaseSmoothForCausalLM
  from .linear import W8A8B8O8Linear, W8A8BFP32O32LinearSiLU, W8A8BFP32OFP32Linear
label_desc: Use image node_id 2 for calling the SmoothLlamaForCausalLM obj's underlying member method code with example usage: obj.get_act_dict(tokenizer, dataset, num_samples, seq_len) and returns: act_dict