(self, prog, mode)
| 1137 | } |
| 1138 | |
| 1139 | def _compile_and_initialize(self, prog, mode): |
| 1140 | compiled_prog = self._compiled_progs.get(mode, None) |
| 1141 | if compiled_prog is not None: |
| 1142 | return compiled_prog |
| 1143 | |
| 1144 | assert self.model._place is not None, ( |
| 1145 | "device is not set, please call `model.prepare()` first" |
| 1146 | ) |
| 1147 | |
| 1148 | place = self.model._place |
| 1149 | |
| 1150 | # XXX *ALL WEIGHTS* should be initialized upon model construction |
| 1151 | # even if `forward()` may run different code path for different mode |
| 1152 | # therefore startup program only needs to run once |
| 1153 | if self._executor is None: |
| 1154 | self._executor = base.Executor(place) |
| 1155 | # XXX incremental initialization |
| 1156 | uninitialized = [] |
| 1157 | for var_py in self._startup_prog.list_vars(): |
| 1158 | var = base.global_scope().find_var(var_py.name) |
| 1159 | if ( |
| 1160 | not var_py.name.startswith('nccl_id') |
| 1161 | and var |
| 1162 | and var.get_tensor()._is_initialized() |
| 1163 | ): |
| 1164 | continue |
| 1165 | |
| 1166 | uninitialized.append(var_py) |
| 1167 | |
| 1168 | # for RawProgramOptimizer, it will insert OP with no outputs like: |
| 1169 | # c_comm_init(inputs={X=['comm_id_0']} |
| 1170 | # but we cannot prune this op. |
| 1171 | block = self._startup_prog.global_block() |
| 1172 | for op in block.ops: |
| 1173 | if op.type == "c_comm_init": |
| 1174 | uninitialized.append(op) |
| 1175 | |
| 1176 | if uninitialized: |
| 1177 | startup_prog = self._startup_prog._prune(uninitialized) |
| 1178 | self._executor.run(startup_prog) |
| 1179 | |
| 1180 | if ( |
| 1181 | self._amp_level == "O2" |
| 1182 | and mode == 'train' |
| 1183 | and core.is_compiled_with_cuda() |
| 1184 | ): |
| 1185 | self.model._optimizer.amp_init(place) |
| 1186 | |
| 1187 | if self._nranks < 2: |
| 1188 | compiled_prog = base.CompiledProgram(prog) |
| 1189 | else: |
| 1190 | compiled_prog = prog |
| 1191 | |
| 1192 | self._compiled_progs[mode] = compiled_prog |
| 1193 | |
| 1194 | |
| 1195 | class DynamicGraphAdapter: |
no test coverage detected