(self)
| 5060 | self.checkpoint_engine.save(state_dict=state, path=save_path) |
| 5061 | |
| 5062 | def _get_buffer_names(self): |
| 5063 | buffer_names = [] |
| 5064 | |
| 5065 | # we save buffer names so that we could extract later the real buffers from the saved |
| 5066 | # state_dict["module"] in the non-zero checkpoint - the buffers are already there but they |
| 5067 | # are intermixed with param placeholders |
| 5068 | |
| 5069 | # have to traverse the tree to be able to skip non-persistent buffers |
| 5070 | def get_layer_named_buffers(module, prefix=""): |
| 5071 | for name, buf in module.named_buffers(recurse=False): |
| 5072 | if buf is not None and name not in module._non_persistent_buffers_set: |
| 5073 | buffer_names.append(prefix + name) |
| 5074 | |
| 5075 | for name, child in module.named_children(): |
| 5076 | if child is not None: |
| 5077 | get_layer_named_buffers(child, prefix + name + ".") |
| 5078 | |
| 5079 | get_layer_named_buffers(self.module, prefix="") |
| 5080 | |
| 5081 | return buffer_names |
| 5082 | |
| 5083 | def _get_param_shape_func(self, param): |
| 5084 | return param.ds_shape if hasattr(param, 'ds_id') else param.shape |
no outgoing calls
no test coverage detected