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 torch_int.nn.bmm import BMM_S8T_S8N_F32T, BMM_S8T_S8N_S8T
  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, repeat_kv, rotate_half
  from transformers.utils import add_start_docstrings_to_model_forward
  from colossalai.inference.tensor_parallel.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