(self, *args, **kwargs)
| 238 | return result |
| 239 | |
| 240 | def __pre_forward_map(self, *args, **kwargs): |
| 241 | # Insert identity op when doing activation checkpointing or pipeline execution. |
| 242 | # Identity op outside activation checkpointing scope will be the endpoint of an activation checkpointing segment. |
| 243 | # Identity op as the first op of a pipeline stage will make backward op depends on the identity op within the stage, |
| 244 | # otherwise the backward op may depends the op in former stage which will make graph creates unnessary buffers. |
| 245 | if self.to(GraphModule)._stage_placement is not None: |
| 246 | |
| 247 | def insert_to_global(t): |
| 248 | assert isinstance(t, Tensor) |
| 249 | return self.__get_or_create_global( |
| 250 | t, self.to(GraphModule)._stage_placement |
| 251 | ) |
| 252 | |
| 253 | args, kwargs = self.__map_io( |
| 254 | "input", insert_to_global, "insert_to_global", *args, **kwargs |
| 255 | ) |
| 256 | |
| 257 | if self.to(GraphModule).activation_checkpointing or ( |
| 258 | self.to(GraphModule).stage_id is not None |
| 259 | and self.to(GraphModule).stage_id >= 0 |
| 260 | ): |
| 261 | |
| 262 | def insert_identity(t): |
| 263 | assert isinstance(t, Tensor) |
| 264 | return self.__get_or_create_identity(t) |
| 265 | |
| 266 | args, kwargs = self.__map_io( |
| 267 | "input", insert_identity, "insert_identity", *args, **kwargs |
| 268 | ) |
| 269 | |
| 270 | return args, kwargs |
| 271 | |
| 272 | def __get_or_create_global(self, input_tensor: Tensor = None, placement=None): |
| 273 | assert input_tensor is not None |
no test coverage detected