Merge pull request #248 from JamesLim-sy/support_O2

Support for BF16-O2

Author: SuperXiang

apps/protein_folding/helixfold/alphafold_paddle/model/all_atom.py

@@ -868,8 +868,8 @@ def between_residue_clash_loss(
     dists_mask *= (residue_index1 < residue_index2)
 
     # Backbone C--N bond between subsequent residues is no clash.
-    c_one_hot = nn.functional.one_hot(paddle.to_tensor([2]), num_classes=14)
-    n_one_hot = nn.functional.one_hot(paddle.to_tensor([0]), num_classes=14)
+    c_one_hot = nn.functional.one_hot(paddle.full(shape=[1], fill_value=2, dtype="int64"), num_classes=14)
+    n_one_hot = nn.functional.one_hot(paddle.full(shape=[1], fill_value=0, dtype="int64"), num_classes=14)
     neighbour_mask = ((residue_index1 + 1) == residue_index2)
     tmp_c_one_hot = paddle.unsqueeze(c_one_hot, axis=[1,2,4])
     tmp_n_one_hot = paddle.unsqueeze(n_one_hot, axis=[1,2,3])
@@ -879,7 +879,7 @@ def between_residue_clash_loss(
 
     # Disulfide bridge between two cysteines is no clash.
     cys_sg_idx = residue_constants.restype_name_to_atom14_names['CYS'].index('SG')
-    cys_sg_one_hot = nn.functional.one_hot(paddle.to_tensor(cys_sg_idx), num_classes=14)
+    cys_sg_one_hot = nn.functional.one_hot(paddle.full(shape=[1], fill_value=cys_sg_idx, dtype="int64"), num_classes=14)
     cys_sg_one_hot1 = paddle.unsqueeze(cys_sg_one_hot, axis=[1,2,4])
     cys_sg_one_hot2 = paddle.unsqueeze(cys_sg_one_hot, axis=[1,2,3])
     disulfide_bonds = (cys_sg_one_hot1 * cys_sg_one_hot2)

apps/protein_folding/helixfold/alphafold_paddle/model/modules.py

@@ -290,6 +290,7 @@ def __init__(self, channel_num, config, global_config):
         }
 
         self.used_heads = []
+        self.heads = []
         for head_name, head_config in sorted(self.config.heads.items()):
             if head_name not in Head_modules:
                 continue
@@ -300,6 +301,7 @@ def __init__(self, channel_num, config, global_config):
 
             head_name_ = Head_names.get(head_name, head_name)
             setattr(self, head_name_, module)
+            self.heads.append(module)
 
     def forward(self,
                 ensembled_batch,
@@ -2173,7 +2175,7 @@ def forward(self, query_embedding, batch, mask_2d):
         Returns:
             A template embedding [N_res, N_res, c_z].
         """
-        assert mask_2d.dtype == query_embedding.dtype
+        assert mask_2d.dtype == query_embedding.dtype, f"mask_2d.dtype ({mask_2d.dtype}) is not the same with query_embedding.dtype ({query_embedding.dtype})!"
         dtype = query_embedding.dtype
         num_res = batch['template_aatype'].shape[1]
         template_mask = batch['template_pseudo_beta_mask']

apps/protein_folding/helixfold/train.py

@@ -30,7 +30,7 @@
 from tensorboardX import SummaryWriter
 
 from utils.utils import get_model_parameter_size, add_to_data_writer, upload_to_hadoop, csv_print
-from utils.utils import get_bf16_op_list
+from utils.utils import get_custom_amp_list
 from utils.metric import ResultsCollect
 from utils.model import RunModel
 from utils.exponential_moving_average import ExponentialMovingAverage, EMA
@@ -53,7 +53,8 @@ def time_me():
     # paddle.device.cuda.synchronize()
     return time.time()
 
-def get_optimizer(opt_config, model):
+
+def get_optimizer(args, opt_config, model):
     if opt_config.grad_clip == 0:
         grad_clip = None
     else:
@@ -74,11 +75,13 @@ def get_optimizer(opt_config, model):
 
     parameters = get_fused_param_groups(model, args.dap_degree > 1 or args.bp_degree > 1)
 
+    multi_precision = (args.precision == "bf16" and args.amp_level == "O2")
     optimizer = paddle.optimizer.Adam(
             learning_rate=lr_scheduler, 
             epsilon=1e-06,
             grad_clip=grad_clip,
-            parameters = parameters
+            parameters=parameters,
+            multi_precision=multi_precision,
         )
     return optimizer, lr_scheduler
 
@@ -141,9 +144,22 @@ def eval(args, model, eval_dataset, compute_loss, cache_dir=None):
             batch['feat'] = align_feat(batch['feat'], args.dap_degree)
             batch['label'] = align_label(batch['label'], args.dap_degree)
 
-        res = model(batch, compute_loss=compute_loss)
+        if args.precision == "bf16" and args.amp_level == "O2":
+            black_list, white_list = get_custom_amp_list()
+            with paddle.amp.auto_cast(enable=True,
+                                      custom_white_list=white_list,
+                                      custom_black_list=black_list,
+                                      level=args.amp_level,
+                                      dtype='bfloat16'):
+                res = model(batch, compute_loss=compute_loss)
+        else:
+            res = model(batch, compute_loss=compute_loss)
         if compute_loss:
             results, loss = res
+            if loss.dtype == paddle.bfloat16:
+                loss = loss.cast("float32").item()
+            else:
+                loss = loss.item()
         else:
             results, loss = res, np.zeros([1])
         s2 = time_me()
@@ -218,8 +234,12 @@ def train(args, cur_step, model, train_data_gen, distill_data_gen, train_config,
     # train
     def _forward_with_precision(batch):
         if args.precision == "bf16":
-            black_list, white_list = get_bf16_op_list()
-            with paddle.amp.auto_cast(level='O1', custom_white_list=white_list, custom_black_list=black_list, dtype='bfloat16'):
+            black_list, white_list = get_custom_amp_list()
+            with paddle.amp.auto_cast(enable=True,
+                                      custom_white_list=white_list,
+                                      custom_black_list=black_list,
+                                      level=args.amp_level,
+                                      dtype='bfloat16'):
                 return model(batch)
         elif args.precision == "fp32":
             return model(batch)
@@ -250,8 +270,7 @@ def _forward_with_precision(batch):
         ema.update()
         optimizer.clear_grad()
 
-    if args.precision == "bf16":
-        loss = loss.cast("float32")
+    loss = loss.cast("float32") if loss.dtype == paddle.bfloat16 else loss
 
     s5 = time_me()
     batch_cost = s5 - s0
@@ -283,6 +302,7 @@ def main(args):
     set_logging_level(args.logging_level)
 
     """main function"""
+    print(f'>>> PaddlePaddle commit: {paddle.version.commit}')
     print(f'>>> args:\n{args}')
     data_config = ml_collections.ConfigDict(json.load(open(args.data_config, 'r')))
     print(f'>>> data_config:\n{data_config}')
@@ -314,7 +334,7 @@ def worker_init_fn(worker_id):
         model_config.model.global_config.dist_model = True
     if args.bp_degree > 1:
         model_config.model.global_config.outer_product_mean_position = 'end'
-    # print(f'>>> model_config:\n{model_config}')
+    print(f'>>> model_config:\n{model_config}')
 
     model = RunModel(train_config, model_config)
 
@@ -377,13 +397,22 @@ def worker_init_fn(worker_id):
 
         model.alphafold.set_state_dict(pd_params)
 
-    optimizer, lr_scheduler = get_optimizer(train_config.optimizer, model)
+    if args.precision == "bf16" and args.amp_level == "O2":
+        print(f"args.amp_level : {args.amp_level}")
+        model = paddle.amp.decorate(
+            models=model,
+            level=args.amp_level,
+            dtype='bfloat16',
+            excluded_layers=model.alphafold.alphafold_iteration.heads
+        )
+
+    optimizer, lr_scheduler = get_optimizer(args, train_config.optimizer, model)
     args.grad_clip = train_config.optimizer.grad_clip
 
     # ema = ExponentialMovingAverage(model, 0.999)
     ema = EMA(optimizer._param_groups, 0.999)
     ema.register()
-    
+
     ### load dataset
     if not args.only_test:
         train_dataset = AF2Dataset(
@@ -473,6 +502,7 @@ def worker_init_fn(worker_id):
     for _ in range(cur_step):
         lr_scheduler.step()
     logging.info('[Main] Start training.')
+
     while True:
         # reset train log info
         if cur_step == 5:
@@ -484,6 +514,7 @@ def worker_init_fn(worker_id):
         # train
         train(args, cur_step, model, train_data_gen, distill_data_gen, train_config, model_config, \
                 lr_scheduler, optimizer, res_collect, train_logger, ema)
+
         if cur_step % args.log_step == 0:
             train_results = res_collect.get_result()
             train_results['lr'] = lr_scheduler.get_lr()
@@ -522,6 +553,7 @@ def worker_init_fn(worker_id):
     parser.add_argument("--model_name", type=str, help='used to choose model config')
     parser.add_argument("--init_model", type=str, default='')
     parser.add_argument("--precision", type=str, choices=['fp32', 'bf16'], default='fp32')
+    parser.add_argument("--amp_level", type=str, default='O1')
     parser.add_argument("--start_step", type=int, default=0)
     parser.add_argument("--train_step", type=int, default=1000)
     parser.add_argument("--batch_size", type=int, default=1)

apps/protein_folding/helixfold/utils/exponential_moving_average.py

@@ -89,8 +89,8 @@ def register(self):
             for p in group['params']:
                 if p.stop_gradient is True:
                     continue
-                self._shadow[id(p)] = paddle.zeros_like(p)
-                self._shadow[id(p)].set_value(p)
+                self._shadow[id(p)] = paddle.zeros_like(p, dtype="float32")
+                self._shadow[id(p)].set_value(p.astype("float32"))
 
     @paddle.no_grad()
     def update(self):
@@ -104,7 +104,7 @@ def update(self):
                     continue
                 new_val = p.detach().clone()
                 old_val = self._shadow[id(p)]
-                new_average = decay * old_val + (1 - decay) * new_val
+                new_average = decay * old_val + (1 - decay) * new_val.astype("float32")
                 self._shadow[id(p)] = new_average
 
         self._update_step += 1
@@ -121,7 +121,10 @@ def apply_shadow(self):
                 assert id(p) in self._shadow
 
                 self._backup[id(p)] = p.detach().clone()
-                p.set_value(self._shadow[id(p)])
+                if p.dtype == paddle.bfloat16:
+                    p.set_value(self._shadow[id(p)].astype(paddle.bfloat16))
+                else:
+                    p.set_value(self._shadow[id(p)])
 
     @paddle.no_grad()
     def restore(self):
@@ -133,4 +136,4 @@ def restore(self):
                     continue
                 assert id(p) in self._backup
                 p.set_value(self._backup[id(p)])
-        self._backup = {}
+        self._backup = {}

apps/protein_folding/helixfold/utils/metric.py

@@ -286,7 +286,15 @@ def get_result(self):
 
     def _extract_loss_dict(self, results):
         """extract value with 'loss' or 'fape' in key"""
+        def _calc_tensor_mean(x):
+            if x.dtype == paddle.bfloat16:
+                x = x.cast("float32")
+            if len(x.shape) == 0:
+                return x.item()
+            else:
+                return x.numpy().mean()
+
         res = tree_flatten(results)
         res = tree_filter(lambda k: 'loss' in k or 'fape' in k, None, res)
-        res = tree_map(lambda x: x.numpy().mean(), res)
+        res = tree_map(lambda x: _calc_tensor_mean(x), res)
         return res

apps/protein_folding/helixfold/utils/utils.py

@@ -19,21 +19,107 @@
 import numpy as np
 import paddle
 
-def get_bf16_op_list():
+
+def get_custom_amp_list():
     """tbd."""
 
     black_list = {"reduce_sum"}
-    white_list = {"concat", "elementwise_add", "elementwise_div", "elementwise_mul", "elementwise_sub", "fill_any_like", "fill_constant", "gather", "gaussian_random",
-        "softmax", "layer_norm", "log_softmax", "matmul_v2", "p_norm", "py_layer", "relu", "scale", "sigmoid", "slice", "softplus", "split", "sqrt", "square", "stack",
-        "sum", "transpose2", "fused_gate_attention", "dropout_nd"}
+    white_list = {
+        "concat",
+        "dropout_nd",
+        "einsum",
+        "elementwise_add",
+        "elementwise_div",
+        "elementwise_mul",
+        "elementwise_sub",
+        "fill_any_like",
+        "fill_constant",
+        "fused_gate_attention",
+        "gather",
+        "gaussian_random",
+        "layer_norm",
+        "log_softmax",
+        "matmul_v2",
+        "p_norm",
+        "py_layer",
+        "relu",
+        "scale",
+        "sigmoid",
+        "slice",
+        "softmax",
+        "softplus",
+        "split",
+        "split_with_num",
+        "sqrt",
+        "square",
+        "stack",
+        "sum",
+        "transpose2",
+        "unsqueeze2",
+        "unstack",
+        "where"
+    }
     return black_list, white_list
 
+
 def get_structure_module_bf16_op_list():
-    black_list = {"reduce_sum", "elementwise_add", "elementwise_div", "elementwise_mul", "elementwise_sub", "fill_any_like", "fill_constant", "gaussian_random", "uniform_random",
-        "softmax", "log_softmax",  "p_norm", "py_layer", "scale", "sigmoid", "softplus", "sqrt", "square", "linspace", "squared_l2_norm", "reduce_mean", "reduce_min", "reduce_prod", "sum", "fused_gate_attention", "dropout_nd", "clip"}
-    white_list = {"layer_norm", "relu", "split", "stack", "gather", "concat", "transpose2", "matmul_v2", "unsqueeze2", "squeeze2", "tile", "slice", "one_hot_v2", "reshape2", "elementwise_max", "elementwise_min", "equal", "greater_than", "less_than", "reduce_max", "eye", "bitwise_or", "abs", "reduce_max", }
+    black_list = {
+        "clip",
+        "dropout_nd",
+        "elementwise_add",
+        "elementwise_div",
+        "elementwise_mul",
+        "elementwise_sub",
+        "fill_any_like",
+        "fill_constant",
+        "fused_gate_attention",
+        "gaussian_random",
+        "linspace",
+        "log_softmax",
+        "p_norm",
+        "py_layer",
+        "reduce_mean",
+        "reduce_min",
+        "reduce_prod",
+        "reduce_sum",
+        "scale",
+        "sigmoid",
+        "softmax",
+        "softplus",
+        "sqrt",
+        "square",
+        "squared_l2_norm",
+        "sum",
+        "uniform_random",
+    }
+    white_list = {
+        "abs",
+        "bitwise_or",
+        "concat",
+        "elementwise_max",
+        "elementwise_min",
+        "equal", 
+        "eye",
+        "gather",
+        "greater_than",
+        "layer_norm",
+        "less_than",
+        "matmul_v2",
+        "one_hot_v2",
+        "reduce_max",
+        "relu",
+        "reshape2",
+        "slice", 
+        "split",
+        "squeeze2",
+        "stack",
+        "transpose2",
+        "unsqueeze2",
+        "tile",
+    }
     return black_list, white_list 
 
+
 def get_model_parameter_size(model):
     """tbd"""
     size = 0
@@ -119,4 +205,4 @@ def csv_print(d):
     keys = sorted(list(d.keys()))
     values = [str(d[k]) for k in keys]
     print(' '.join([str(x) for x in keys]))
-    print(' '.join([str(x) for x in values]))
+    print(' '.join([str(x) for x in values]))