1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

def construct_stages(self, submeshes: List[Tuple], stages_ids: List[Tuple]):
    # construct layers for each stage
    first_part, module_list, last_part = self.parse_func(self.model)
    modules = list()
    num = len(stages_ids)
    for idx in range(num):
        submesh = submeshes[idx]
        stage_id = stages_ids[idx]
        # get stage layers from user config stage ids in module list
        layers = list(module_list[stage_id[0]: stage_id[1] + 1])
        if idx == 0 and first_part is not None:
            # concat module first part(if exists) bef module list to stage_0
            layers = first_part + layers
        if idx == num - 1 and last_part is not None:
            # concat module last part(if exists) aft module list to last stage
            layers.extend(last_part)

        modules.append(layers)
        # deepcopy module for xla device tracing use
        stage_module = [copy.deepcopy(layer) for layer in layers]
        self.stages.append(
            Stage(idx, stage_module, submesh, self, ))
    return modules

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

def shard_sora_one_stage(modules, shard_strategy, mesh):
    total_len = len(modules)
    # first 5 modules are embedding layers
    for i in range(0, 5):
        shard_sora_embedding(modules[i], shard_strategy, mesh)
    for i in range(5, total_len - 2):
        shard_sora_block(modules[i][0], shard_strategy, mesh)  # shard spatial
        shard_sora_block(modules[i][1], shard_strategy, mesh)  # shard temporal
    shard_sora_final(modules[-1], shard_strategy, mesh)


def shard_sora_first_stage(modules, shard_strategy, mesh):
    for i in range(0, 5):
        shard_sora_embedding(modules[i], shard_strategy, mesh)
    for i in range(5, len(modules)):
        shard_sora_block(modules[i][0], shard_strategy, mesh)  # shard spatial
        shard_sora_block(modules[i][1], shard_strategy, mesh)  # shard temporal


def shard_sora_stage(modules, shard_strategy, mesh):
    for module in modules:
        shard_sora_block(module[0], shard_strategy, mesh)  # shard spatial
        shard_sora_block(module[1], shard_strategy, mesh)  # shard temporal


def shard_sora_last_stage(modules, shard_strategy, mesh):
    total_len = len(modules)
    for i in range(0, total_len - 2):
        shard_sora_block(modules[i][0], shard_strategy, mesh)  # shard spatial
        shard_sora_block(modules[i][1], shard_strategy, mesh)  # shard temporal
    # skip norm layer mark sharding
    shard_sora_final(modules[total_len - 1], shard_strategy, mesh)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59

def trace_stage_mhlo_graph(self, check_res=False):
    """
    trace stage nn modules to mhlo graph
    """
    # (NOTE): construct xla mesh before trace tensors generate,
    # i.e., before any xla device call to avoid xla computation client construct
    xla_mesh = None
    if self.shard_func is not None:
        xla_mesh = self._construct_stage_xla_mesh()  # create mesh from submesh info
    # Create xla device trace tensors, move module to xla device
    if self.stage_id == 0:
        self.trace_tensors = self._generate_trace_tensors()
    else:
        z = self.parent_pipeline.stages[self.stage_id -1].outputs
        y = self.parent_pipeline.stages[0].y_embedded.to('cpu').to(xm.xla_device())
        t_mlp = self.parent_pipeline.stages[0].t_mlp.to('cpu').to(xm.xla_device())
        self.trace_tensors = [z, y, t_mlp]
    for module in self.modules:
        if isinstance(module, tuple):
            for mod in module:
                mod.to('cpu').to(xm.xla_device())  # first load to cpu
        else:
            module.to('cpu').to(xm.xla_device())
    # get pipeline exec mode
    assert self.parent_pipeline is not None
    exec_mode = self.parent_pipeline.exec_mode
    # load lora cofifg
    lora_config = self.parent_pipeline.lora_config

    print("Enter trace mhlo graph for stage: ", self.stage_id)
    # Trigger shard func to mark sharding the model
    if self.shard_func is not None:
        self.shard_func(self.modules, self.shard_strategy, xla_mesh)

    if exec_mode == EXEC_MODE.INFER:
        # set stage name & dump file path
        self._set_stage_name_dump_file(
            exec_mode, "fw")
        num_sampling_steps = 30
        num_timesteps = 1000
        timesteps = [(1.0 - i / num_sampling_steps) * num_timesteps for i in range(num_sampling_steps)]
        # FIXME: 原先是为每个stage单独生成trace_tensor, 现在要把上一个的结果传给下一个 stage
        #for i in range(30):
        start = sum(self.parent_pipeline.pipeline_patches_height_list[:self.stage_id - 1]) if self.stage_id != 0 else 0
        end = start + self.parent_pipeline.pipeline_patches_height_list[self.stage_id]
    
        if self.stage_id == 0:
            outputs = self._forward([self.trace_tensors[0][...,start:end,:]] + self.trace_tensors[1:], xla_mesh)  # outputs is a list
        else:
            outputs = self._forward(self.trace_tensors, xla_mesh)

        if check_res:
            # check xla results compared to gpu results
            check_result_error(self.outputs, outputs)
        else:
            # use torch xla _get_xla_tensors_hlo interface
            # to eliminate redundant live tensors as ret values
            os.environ["XLA_DUMP_POST_OPTIMIZATIONS"] = "true"
            torch_xla._XLAC._get_xla_tensors_hlo(outputs)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

def collect_rms_ops(self):
  rms_collector = RMSCollector()
  rms_collector.visit(self.root_op)
  self.rms_locs = rms_collector.rms_locs
  # construct attention block & ffn block ranges
  # exclude the rsqrt in T2IFinalLayer
  att_rm_locs = self.rms_locs if len(self.rms_locs) % 2 == 0 else self.rms_locs[:-1] 

  for i in range(8, len(att_rm_locs), 4):  # a block has 4 rsqrt, start from 2nd block to avoid embedding
      self.spt_qkv_ranges.append((att_rm_locs[i+0], att_rm_locs[i+1]))
      self.spt_attn_ranges.append((att_rm_locs[i+2], att_rm_locs[i+3]))
      self.cro_attn_ranges.append((att_rm_locs[i+2], att_rm_locs[i+3]))

  for i in range(8, len(att_rm_locs), 4):  # ORG: range(8, len(att_rm_locs), 4): 
      start = self.rms_locs[i+3]
      if i+4 >= len(self.rms_locs):
          end = None
      else:
          end = self.rms_locs[i+4]
      self.ffn_ranges.append((start, end))

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47

def visit_dot(self, node):
    dot_lineno = _parse_loc_lineno(node)

    if self.block_cnt < len(self.spt_attn_ranges):
        spt_att_range = self.spt_attn_ranges[self.block_cnt]
        cro_att_range = self.cro_attn_ranges[self.block_cnt]
        spt_qkv_range = self.spt_qkv_ranges[self.block_cnt]
        ffn_range = self.ffn_ranges[self.block_cnt]

        # lie in RMS ops closed attention block
        if dot_lineno > spt_att_range[0] and dot_lineno < spt_att_range[1]:
            #import pdb;pdb.set_trace()
            self.att_block_dots.append(node)
            self.spt_dot_cnt += 1
        elif dot_lineno > cro_att_range[0] and dot_lineno < cro_att_range[1]:
            self.att_block_dots.append(node)
            self.cro_att_dot_cnt += 1
        # lie ffn block
        if dot_lineno > spt_qkv_range[0] and dot_lineno < spt_qkv_range[1]:
            self.spt_qkv_cnt += 1
            self.ffn_block_dots.append(node)
            # pixart pass
        elif dot_lineno > ffn_range[0]:
            if ffn_range[1] is not None:
                if dot_lineno < ffn_range[1]:
                    self.ffn_block_dots.append(node)
                    self.ffn_dot_cnt += 1
            else:
                if self.ffn_dot_cnt < 2:                 
                    self.ffn_block_dots.append(node)
                    self.ffn_dot_cnt += 1
        # Traversal of one block
        if self.spt_qkv_cnt == 1 and self.spt_att_dot_cnt == 4 and \
            self.spt_dot_cnt == 4 and self.ffn_dot_cnt == 2:
            self.attention_blocks.append(self.att_block_dots)
            self.ffn_blocks.append(self.ffn_block_dots)
            self.block_cnt += 1
            # reset each block level counters
            self.spt_qkv_cnt = 0
            self.spt_att_dot_cnt = 0
            self.spt_dot_cnt = 0
            self.ffn_dot_cnt = 0
        
            self.att_block_dots = []
            self.ffn_block_dots = []

    self.generic_visit(node)

1
2
3
4
5
6
7

spatial_config = {"B": self.config["B_spt"], "S_Q": self.config["S_Q_spt"], "S_KV": self.config["S_KV_spt"], "D_QKV": self.config["D_QKV"], 
                  "H_QKV": self.config["H_QKV"], "N_A": self.config["N_A"], "H_A": self.config["H_A"], "D_O": self.config["D_O_spt"], "H_O": self.config["H_O_spt"] }
temporal_config = {"B": self.config["B_tmp"], "S_Q": self.config["S_Q_tmp"], "S_KV": self.config["S_KV_tmp"], "D_QKV": self.config["D_QKV"], 
                  "H_QKV": self.config["H_QKV"], "N_A": self.config["N_A"], "H_A": self.config["H_A"], "D_O": self.config["D_O_tmp"], "H_O": self.config["H_O_tmp"] }
cross_config = {"B": self.config["B_cro"], "S_Q": self.config["S_Q_cro"], "S_KV": self.config["S_KV_cro"], "D_QKV": self.config["D_QKV"], 
                "H_QKV": self.config["H_QKV"],"N_A": self.config["N_A"], "H_A": self.config["H_A"], "D_O": self.config["D_O_cro"], "H_O": self.config["H_O_cro"],
                "D_FU": self.config["D_FU"], "H_FU": self.config["H_FU"], "D_FD": self.config["D_FD"], "H_FD": self.config["H_FD"]}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

GB = 2**30

class Gate_ResAdd():
  '''
  Construct each op after MHSA on the config file
  '''
  def __init__(self, config: dict, name: str) -> None:
      self.config = config
      self.name = name
      # {name:{type:"", size:"", ishape:[], wshape:[]/None, oshape:[]}}
      self.ops = {}
      self.construct_model()
  
  def construct_model(self):
      GB = 2**30
      ResAdd_input_shape = [self.config['B'], self.config['S_Q'], self.config['D_O']]
      ResAdd_weight_shape = [1, self.config['D_O']]
      ResAdd_output_shape = ResAdd_input_shape
      ResAdd_compute = 2*ResAdd_input_shape[0]*ResAdd_input_shape[1]*ResAdd_input_shape[2]/GB
      self.ops[self.name+"_"+"ResAdd"] = {"name":"ResAdd", 
                                          "type": "Vector", 
                                          "ishape": ResAdd_input_shape, 
                                          "wshape": ResAdd_weight_shape, 
                                          "oshape": ResAdd_output_shape, 
                                          "compute": ResAdd_compute}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

class STDIT2_block():
  def __init__(self, config) -> None:
      self.config = config
      # {name:{type:"", size:"", ishape:[], wshape:[]/None, oshape:[]}}
      self.ops = {}
      self.construct_model()
  
  def construct_model(self):
      spatial_config = {"B": self.config["B_spt"], "S_Q": self.config["S_Q_spt"], "S_KV": self.config["S_KV_spt"], "D_QKV": self.config["D_QKV"], 
                      "H_QKV": self.config["H_QKV"], "N_A": self.config["N_A"], "H_A": self.config["H_A"], "D_O": self.config["D_O_spt"], "H_O": self.config["H_O_spt"] }
      temporal_config = {"B": self.config["B_tmp"], "S_Q": self.config["S_Q_tmp"], "S_KV": self.config["S_KV_tmp"], "D_QKV": self.config["D_QKV"], 
                      "H_QKV": self.config["H_QKV"], "N_A": self.config["N_A"], "H_A": self.config["H_A"], "D_O": self.config["D_O_tmp"], "H_O": self.config["H_O_tmp"] }
      cross_config = {"B": self.config["B_cro"], "S_Q": self.config["S_Q_cro"], "S_KV": self.config["S_KV_cro"], "D_QKV": self.config["D_QKV"], 
                      "H_QKV": self.config["H_QKV"],"N_A": self.config["N_A"], "H_A": self.config["H_A"], "D_O": self.config["D_O_cro"], "H_O": self.config["H_O_cro"],
                      "D_FU": self.config["D_FU"], "H_FU": self.config["H_FU"], "D_FD": self.config["D_FD"], "H_FD": self.config["H_FD"]}

      self.spatial_modulate = Modulate(spatial_config, name="spatial")
      self.spatial_block = MHSA_block(spatial_config, name="spatial")
      self.spatial_gate_resadd = Gate_ResAdd(spatial_config, name="spatial")
      self.temporal_modulate = Modulate(temporal_config, name="temporal")
      self.temporal_block = MHSA_block(temporal_config, name="temporal")
      self.temporal_gate_resadd = Gate_ResAdd(temporal_config, name="temporal")
      self.cross_block = MHSA_block(cross_config, name="cross")
      self.cross_gate_resadd = Gate_ResAdd(cross_config, name="cross")
      self.mlp_modulate = Modulate(cross_config, name="mlp")
      self.ffn_block = FFN_block(cross_config)
      self.mlp_gate_resadd = Gate_ResAdd(cross_config, name="mlp")

      op_list = [self.spatial_modulate.ops, self.spatial_block.ops, self.spatial_gate_resadd.ops, 
                  self.temporal_modulate.ops, self.temporal_block.ops, self.temporal_gate_resadd.ops, 
                  self.cross_block.ops, self.cross_gate_resadd.ops, self.mlp_modulate.ops, self.ffn_block.ops, self.mlp_gate_resadd.ops]

      for op_dict in op_list:
          self.ops.update(op_dict)
      print(self.ops.keys())

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65

def STDIT2_mapper(model, arch, QKV_fusion=True, preset=True, details=True):
  config = model.config
  Layers = config['L']
  spatial_config = {'B': config['B_spt'], 'S_Q': config['S_Q_spt'], 'S_KV': config['S_KV_spt'], 'H_A': config['H_A'], 'N_A': config['N_A'], 'Q': config['Q']}
  temporal_config = {'B': config['B_tmp'], 'S_Q': config['S_Q_tmp'], 'S_KV': config['S_KV_tmp'], 'H_A': config['H_A'], 'N_A': config['N_A'], 'Q': config['Q']}
  cross_config = {'B': config['B_cro'], 'S_Q': config['S_Q_cro'], 'S_KV': config['S_KV_cro'], 'H_A': config['H_A'], 'N_A': config['N_A'], 'Q': config['Q']}
  ops = model.ops
  mapping_result = {}
  
  '''=========================
  == Spatial Branch Mapping ==
  ========================='''
  TmTn = [256, 32] if preset else None
  mapping_result['spatial_Modulate'] = vector_mapper(ops['spatial_Modulate'],arch,splits=None,details=details)
  mapping_result['spatial_RMSNorm']= vector_mapper(ops['spatial_RMSNorm'],arch,splits=None,details=details)
  mapping_result['spatial_Q_proj'] = gemm_auto_opt_mapper(ops['spatial_Q_proj'], arch, TmTn=TmTn, details=details)
  mapping_result['spatial_K_proj'] = gemm_auto_opt_mapper(ops['spatial_K_proj'], arch, TmTn=TmTn, details=details)
  mapping_result['spatial_V_proj'] = gemm_auto_opt_mapper(ops['spatial_V_proj'], arch, TmTn=TmTn, details=details)
  Tx_Ty = [256, 256] if preset else None
  mapping_result['spatial_Flashatten'] = flashatten_mapper(spatial_config, arch, Tx_Ty=Tx_Ty, details=details, Head_fused=True)  # FIXME
  mapping_result['spatial_ResAdd']=vector_mapper(ops['spatial_ResAdd'],arch,splits=None,details=details)
  
  '''==========================
  == Temporal Branch Mapping ==
  =========================='''
  mapping_result['temporal_Modulate'] = vector_mapper(ops['temporal_Modulate'],arch,splits=None,details=details)  # 切分 30 份也无法满足SRAM要求
  mapping_result['temporal_RMSNorm']= vector_mapper(ops['temporal_RMSNorm'],arch,splits=None,details=details)
  mapping_result['temporal_Q_proj'] = gemm_auto_opt_mapper(ops['temporal_Q_proj'], arch, TmTn=TmTn, details=details)
  mapping_result['temporal_K_proj'] = gemm_auto_opt_mapper(ops['temporal_K_proj'], arch, TmTn=TmTn, details=details)
  mapping_result['temporal_V_proj'] = gemm_auto_opt_mapper(ops['temporal_V_proj'], arch, TmTn=TmTn, details=details)
  Tx_Ty = [256, 256] if preset else None
  mapping_result['temporal_Flashatten'] = flashatten_mapper(temporal_config, arch, Tx_Ty=Tx_Ty, details=details, Head_fused=True)  # FIXME
  mapping_result['temporal_ResAdd']=vector_mapper(ops['temporal_ResAdd'],arch,splits=None,details=details)
  
  '''====================================
  == Cross Branch Mapping 2x per block ==
  ===================================='''
  #mapping_result['spatial_RMSNorm']= vector_mapper(ops['spatial_RMSNorm'],arch,splits=None,details=details)
  mapping_result['cross_Q_proj'] =  gemm_auto_opt_mapper(ops['cross_Q_proj'], arch, TmTn=TmTn, details=details)
  mapping_result['cross_Q_proj_2'] = mapping_result['cross_Q_proj']
  mapping_result['cross_K_proj'] =  gemm_auto_opt_mapper(ops['cross_K_proj'], arch, TmTn=TmTn, details=details)
  mapping_result['cross_K_proj_2'] = mapping_result['cross_K_proj']
  mapping_result['cross_V_proj'] =  gemm_auto_opt_mapper(ops['cross_V_proj'], arch, TmTn=TmTn, details=details)
  mapping_result['cross_V_proj_2'] = mapping_result['cross_V_proj']
  Tx_Ty = [256, 256] if preset else None
  mapping_result['cross_Flashatten'] =  flashatten_mapper(cross_config, arch, Tx_Ty=Tx_Ty, details=details, Head_fused=True)  # FIXME
  mapping_result['cross_Flashatten_2'] = mapping_result['cross_Flashatten']
  mapping_result['cross_ResAdd'] =  vector_mapper(ops['cross_ResAdd'],arch,splits=None,details=details)  
  # HACK: Gate_ResAdd *2 了, cross 无gate 这里无 _2
  
  
  '''====================================
  == Feed Forward Network 2x per block ==
  ===================================='''
  mapping_result['mlp_Modulate'] = vector_mapper(ops['mlp_Modulate'],arch,splits=None,details=details)
  mapping_result['mlp_Modulate_2'] = mapping_result['mlp_Modulate']
  mapping_result['FFNup&SiLU'] = gemm_auto_opt_mapper(ops['FFNup'],arch,TmTn=TmTn,fusion_op2=ops['SiLU'],details=details)
  mapping_result['FFNup&SiLU_2'] = mapping_result['FFNup&SiLU']
  # mapping_result['FFNgate'] = gemm_auto_opt_mapper(ops['FFNgate'], arch, TmTn=TmTn, details=details)
  # mapping_result['Hadamard'] = vector_mapper(ops['Hadamard'], arch, splits=None)
  TmTn = [4, 128] if preset else None
  mapping_result['FFNdown'] = gemm_auto_opt_mapper(ops['FFNdown'], arch, TmTn=TmTn, details=details)
  mapping_result['FFNdown_2'] = mapping_result['FFNdown']
  mapping_result['mlp_ResAdd'] = vector_mapper(ops['mlp_ResAdd'], arch, splits=None, details=details)
  mapping_result['mlp_ResAdd_2'] = mapping_result['mlp_ResAdd']