r"""Set parallel context
(args_opt)
| 102 | |
| 103 | |
| 104 | def set_parallel_context(args_opt): |
| 105 | r"""Set parallel context""" |
| 106 | D.init() |
| 107 | device_num = D.get_group_size() |
| 108 | rank = D.get_rank() |
| 109 | print("rank_id is {}, device_num is {}".format(rank, device_num)) |
| 110 | context.reset_auto_parallel_context() |
| 111 | context.set_auto_parallel_context( |
| 112 | parallel_mode=ParallelMode.SEMI_AUTO_PARALLEL, gradients_mean=False, |
| 113 | full_batch=bool(args_opt.full_batch), strategy_ckpt_load_file=args_opt.strategy_load_ckpt_path, |
| 114 | enable_parallel_optimizer=bool(args_opt.optimizer_shard), strategy_ckpt_save_file='strategy.ckpt', |
| 115 | optimizer_weight_shard_size=16) |
| 116 | set_algo_parameters(elementwise_op_strategy_follow=True) |
| 117 | _set_multi_subgraphs() |
| 118 | return rank, device_num |
| 119 | |
| 120 | |
| 121 | def run_train(args_opt): |