Alan Liu commited on
Commit
79bea76
1 Parent(s): 5607124

fix len(modules)=0 case

Browse files
Files changed (1) hide show
  1. calc_util.py +63 -55
calc_util.py CHANGED
@@ -65,10 +65,11 @@ def word_embedding_operation(model_config, inference_config):
65
  #\end{equation}
66
  if model_config['module_classes']:
67
  modules = get_module_tensors_matched(lambda x: 'embed' in x and 'pos' not in x, model_config['module_classes'])
68
- A = [inference_config['batchsize'], inference_config['input_seq_length'], modules[0][0]]
69
- B = modules[0]
70
- op_count = matrix_operation(A, B)
71
- return op_count
 
72
 
73
  A = [inference_config['batchsize'], inference_config['input_seq_length'], model_config['vocab_size']]
74
  B = [model_config['vocab_size'], model_config['hidden_size']]
@@ -79,7 +80,8 @@ def word_embedding_operation(model_config, inference_config):
79
  def positional_embedding_operation(model_config, inference_config):
80
  if model_config['module_classes']:
81
  modules = get_module_tensors_matched(lambda x: 'embed' in x and 'pos' in x, model_config['module_classes'])
82
- return multiplication_in_int64([inference_config['batchsize'], inference_config['input_seq_length'], modules[0][-1]])
 
83
 
84
  return multiplication_in_int64([inference_config['batchsize'], inference_config['input_seq_length'], model_config['hidden_size']])
85
 
@@ -87,15 +89,16 @@ def positional_embedding_operation(model_config, inference_config):
87
  def attention_K_operation(model_config, inference_config, seq_length):
88
  if model_config['module_classes']:
89
  modules = get_module_tensors_matched(lambda x: 'att' in x and 'k' in x , model_config['module_classes'])
90
- total = 0
91
- for module in modules:
92
- if len(module) > 1:
93
- A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
94
- B = [model_config['hidden_size'], model_config['hidden_size_per_head']]
95
- total += model_config['num_attention_heads']*matrix_operation(A, B)
96
- else:
97
- total += model_config['hidden_size']
98
- return total
 
99
 
100
  A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
101
  B = [model_config['hidden_size'], model_config['hidden_size_per_head']]
@@ -104,15 +107,16 @@ def attention_K_operation(model_config, inference_config, seq_length):
104
  def attention_Q_operation(model_config, inference_config, seq_length):
105
  if model_config['module_classes']:
106
  modules = get_module_tensors_matched(lambda x: 'att' in x and 'q' in x , model_config['module_classes'])
107
- total = 0
108
- for module in modules:
109
- if len(module) > 1:
110
- A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
111
- B = [model_config['hidden_size'], model_config['hidden_size_per_head']]
112
- total += model_config['num_attention_heads']*matrix_operation(A, B)
113
- else:
114
- total += model_config['hidden_size']
115
- return total
 
116
 
117
  A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
118
  B = [model_config['hidden_size'], model_config['hidden_size_per_head']]
@@ -121,15 +125,16 @@ def attention_Q_operation(model_config, inference_config, seq_length):
121
  def attention_V_operation(model_config, inference_config, seq_length):
122
  if model_config['module_classes']:
123
  modules = get_module_tensors_matched(lambda x: 'att' in x and 'v' in x , model_config['module_classes'])
124
- total = 0
125
- for module in modules:
126
- if len(module) > 1:
127
- A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
128
- B = [model_config['hidden_size'], model_config['hidden_size_per_head']]
129
- total += model_config['num_attention_heads']*matrix_operation(A, B)
130
- else:
131
- total += model_config['hidden_size']
132
- return total
 
133
 
134
  A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
135
  B = [model_config['hidden_size'], model_config['hidden_size_per_head']]
@@ -155,15 +160,16 @@ def attention_multV_operation(model_config, inference_config, seq_length_Q, seq_
155
  def attention_out_operation(model_config, inference_config, seq_length):
156
  if model_config['module_classes']:
157
  modules = get_module_tensors_matched(lambda x: 'att' in x and 'k' in x , model_config['module_classes'])
158
- total = 0
159
- for module in modules:
160
- if len(module) > 1:
161
- A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
162
- B = [model_config['hidden_size'], model_config['hidden_size']]
163
- total += matrix_operation(A, B)
164
- else:
165
- total += model_config['hidden_size']
166
- return total
 
167
 
168
  A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
169
  B = [model_config['hidden_size'], model_config['hidden_size']]
@@ -174,11 +180,12 @@ def layernorm_operation(model_config, inference_config, seq_length):
174
  # 5 is a modeled value
175
  if model_config['module_classes']:
176
  modules = get_module_tensors_matched(lambda x: 'norm' in x, model_config['module_classes'])
177
- total = 0
178
- for module in modules:
179
- total += model_config['hidden_size']
180
- return 5*total
181
-
 
182
  layernorm_operation = (5*inference_config['batchsize']*seq_length*model_config['hidden_size'])
183
  return model_config['num_hidden_layers'] * model_config['layernorm_operation'] * layernorm_operation
184
 
@@ -186,15 +193,16 @@ def layernorm_operation(model_config, inference_config, seq_length):
186
  def mlp_operation(model_config, inference_config, seq_length):
187
  if model_config['module_classes']:
188
  modules = get_module_tensors_matched(lambda x: 'fc' in x or 'mlp' in x, model_config['module_classes'])
189
- total = 0
190
- for module in modules:
191
- if len(module) > 1:
192
- A = [inference_config['batchsize'], seq_length, module[1]]
193
- B = [module[1], module[0]]
194
- total += matrix_operation(A, B)
195
- else:
196
- total += modules[-1][0]
197
- return total
 
198
 
199
  A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
200
  B = [model_config['hidden_size'], model_config['intermediate_size']]
 
65
  #\end{equation}
66
  if model_config['module_classes']:
67
  modules = get_module_tensors_matched(lambda x: 'embed' in x and 'pos' not in x, model_config['module_classes'])
68
+ if len(modules) > 0:
69
+ A = [inference_config['batchsize'], inference_config['input_seq_length'], modules[0][0]]
70
+ B = modules[0]
71
+ op_count = matrix_operation(A, B)
72
+ return op_count
73
 
74
  A = [inference_config['batchsize'], inference_config['input_seq_length'], model_config['vocab_size']]
75
  B = [model_config['vocab_size'], model_config['hidden_size']]
 
80
  def positional_embedding_operation(model_config, inference_config):
81
  if model_config['module_classes']:
82
  modules = get_module_tensors_matched(lambda x: 'embed' in x and 'pos' in x, model_config['module_classes'])
83
+ if len(modules) > 0:
84
+ return multiplication_in_int64([inference_config['batchsize'], inference_config['input_seq_length'], modules[0][-1]])
85
 
86
  return multiplication_in_int64([inference_config['batchsize'], inference_config['input_seq_length'], model_config['hidden_size']])
87
 
 
89
  def attention_K_operation(model_config, inference_config, seq_length):
90
  if model_config['module_classes']:
91
  modules = get_module_tensors_matched(lambda x: 'att' in x and 'k' in x , model_config['module_classes'])
92
+ if len(modules) > 0:
93
+ total = 0
94
+ for module in modules:
95
+ if len(module) > 1:
96
+ A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
97
+ B = [model_config['hidden_size'], model_config['hidden_size_per_head']]
98
+ total += model_config['num_attention_heads']*matrix_operation(A, B)
99
+ else:
100
+ total += model_config['hidden_size']
101
+ return total
102
 
103
  A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
104
  B = [model_config['hidden_size'], model_config['hidden_size_per_head']]
 
107
  def attention_Q_operation(model_config, inference_config, seq_length):
108
  if model_config['module_classes']:
109
  modules = get_module_tensors_matched(lambda x: 'att' in x and 'q' in x , model_config['module_classes'])
110
+ if len(modules) > 0:
111
+ total = 0
112
+ for module in modules:
113
+ if len(module) > 1:
114
+ A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
115
+ B = [model_config['hidden_size'], model_config['hidden_size_per_head']]
116
+ total += model_config['num_attention_heads']*matrix_operation(A, B)
117
+ else:
118
+ total += model_config['hidden_size']
119
+ return total
120
 
121
  A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
122
  B = [model_config['hidden_size'], model_config['hidden_size_per_head']]
 
125
  def attention_V_operation(model_config, inference_config, seq_length):
126
  if model_config['module_classes']:
127
  modules = get_module_tensors_matched(lambda x: 'att' in x and 'v' in x , model_config['module_classes'])
128
+ if len(modules) > 0:
129
+ total = 0
130
+ for module in modules:
131
+ if len(module) > 1:
132
+ A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
133
+ B = [model_config['hidden_size'], model_config['hidden_size_per_head']]
134
+ total += model_config['num_attention_heads']*matrix_operation(A, B)
135
+ else:
136
+ total += model_config['hidden_size']
137
+ return total
138
 
139
  A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
140
  B = [model_config['hidden_size'], model_config['hidden_size_per_head']]
 
160
  def attention_out_operation(model_config, inference_config, seq_length):
161
  if model_config['module_classes']:
162
  modules = get_module_tensors_matched(lambda x: 'att' in x and 'k' in x , model_config['module_classes'])
163
+ if len(modules) > 0:
164
+ total = 0
165
+ for module in modules:
166
+ if len(module) > 1:
167
+ A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
168
+ B = [model_config['hidden_size'], model_config['hidden_size']]
169
+ total += matrix_operation(A, B)
170
+ else:
171
+ total += model_config['hidden_size']
172
+ return total
173
 
174
  A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
175
  B = [model_config['hidden_size'], model_config['hidden_size']]
 
180
  # 5 is a modeled value
181
  if model_config['module_classes']:
182
  modules = get_module_tensors_matched(lambda x: 'norm' in x, model_config['module_classes'])
183
+ if len(modules) > 0:
184
+ total = 0
185
+ for module in modules:
186
+ total += model_config['hidden_size']
187
+ return 5*total
188
+
189
  layernorm_operation = (5*inference_config['batchsize']*seq_length*model_config['hidden_size'])
190
  return model_config['num_hidden_layers'] * model_config['layernorm_operation'] * layernorm_operation
191
 
 
193
  def mlp_operation(model_config, inference_config, seq_length):
194
  if model_config['module_classes']:
195
  modules = get_module_tensors_matched(lambda x: 'fc' in x or 'mlp' in x, model_config['module_classes'])
196
+ if len(modules) > 0:
197
+ total = 0
198
+ for module in modules:
199
+ if len(module) > 1:
200
+ A = [inference_config['batchsize'], seq_length, module[1]]
201
+ B = [module[1], module[0]]
202
+ total += matrix_operation(A, B)
203
+ else:
204
+ total += modules[-1][0]
205
+ return total
206
 
207
  A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
208
  B = [model_config['hidden_size'], model_config['intermediate_size']]