MCPcopy Index your code
hub / github.com/zai-org/CodeGeeX / run_train

Function run_train

codegeex/mindspore/train.py:121–330  ·  view source on GitHub ↗

r"""The main training process.

(args_opt)

Source from the content-addressed store, hash-verified

119
120
121def run_train(args_opt):
122 r"""The main training process."""
123 os.environ["HCCL_CONNECT_TIMEOUT"] = "2000"
124 # Set execution mode
125 context.set_context(
126 mode=context.GRAPH_MODE, device_target=args_opt.device_target
127 )
128 if args_opt.profiling:
129 profiler = Profiler(output_path="/cache/profiler_data")
130 context.set_context(variable_memory_max_size="30GB")
131 # Set parallel context
132 rank = 0
133 device_num = 1
134 if args_opt.distribute == "true":
135 rank, device_num = set_parallel_context(args_opt)
136 context.set_context(
137 save_graphs=False,
138 save_graphs_path="/cache/graphs_of_device_id_" + str(rank),
139 )
140 cache_url = '/cache/Data/'
141 eval_cache_url = '/cache/EvalData/'
142 if not args_opt.offline:
143 download_data(src_data_url=args_opt.data_url, tgt_data_path=cache_url, rank=rank)
144 download_data(src_data_url=args_opt.eval_data_url, tgt_data_path=eval_cache_url, rank=rank)
145 # Set model property
146 model_parallel_num = args_opt.op_level_model_parallel_num
147 data_parallel_num = int(device_num / model_parallel_num)
148 batch_size = args_opt.per_batch_size * data_parallel_num
149 parallel_config = TransformerOpParallelConfig(data_parallel=data_parallel_num, model_parallel=model_parallel_num,
150 pipeline_stage=args_opt.stage_num,
151 micro_batch_num=args_opt.micro_size,
152 optimizer_shard=bool(args_opt.optimizer_shard),
153 vocab_emb_dp=bool(args_opt.word_emb_dp), recompute=True,
154 gradient_aggregation_group=args_opt.gradient_aggregation_group)
155
156 micro_interleaved_size = args_opt.micro_interleaved_size
157 config = PanguAlphaConfig(
158 batch_size=batch_size // micro_interleaved_size,
159 num_heads=args_opt.num_heads,
160 hidden_size=args_opt.embedding_size,
161 seq_length=args_opt.seq_length,
162 vocab_size=args_opt.vocab_size,
163 num_layers=args_opt.num_layers,
164 ffn_hidden_size=args_opt.embedding_size * 4,
165 eod_token=args_opt.eod_id,
166 load_ckpt_path=args_opt.load_ckpt_path,
167 param_init_type=mstype.float32
168 if args_opt.param_init_type == "fp32"
169 else mstype.float16,
170 dropout_rate=args_opt.dropout_rate,
171 enable_offload=bool(args_opt.opt_offload),
172 use_moe=bool(args_opt.use_moe),
173 per_dp_dim_expert_num=args_opt.per_dp_dim_expert_num,
174 hidden_act="fast_gelu" if args_opt.device_target != "GPU" else "gelu",
175 parallel_config=parallel_config,
176 )
177 print("===config is: ", config, flush=True)
178 # Define network

Callers 1

train.pyFile · 0.70

Calls 15

download_dataFunction · 0.90
PanguAlphaConfigClass · 0.90
PanguAlphaModelClass · 0.90
PanGUAlphaWithLossClass · 0.90
LearningRateClass · 0.90
AdamWeightDecayOpClass · 0.90
create_datasetFunction · 0.90
PPLMetricClass · 0.90
ValidationLossClass · 0.90
EvalCallBackClass · 0.90

Tested by

no test coverage detected