| from contextlib import contextmanager |
| from typing import * |
| import math |
| from ..modules import sparse as sp |
| from ..utils.elastic_utils import ElasticModuleMixin |
|
|
|
|
| class SparseTransformerElasticMixin(ElasticModuleMixin): |
| def _get_input_size(self, x: sp.SparseTensor, *args, **kwargs): |
| return x.feats.shape[0] |
| |
| @contextmanager |
| def with_mem_ratio(self, mem_ratio=1.0): |
| if mem_ratio == 1.0: |
| yield 1.0 |
| return |
| num_blocks = len(self.blocks) |
| num_checkpoint_blocks = min(math.ceil((1 - mem_ratio) * num_blocks) + 1, num_blocks) |
| exact_mem_ratio = 1 - (num_checkpoint_blocks - 1) / num_blocks |
| for i in range(num_blocks): |
| self.blocks[i].use_checkpoint = i < num_checkpoint_blocks |
| yield exact_mem_ratio |
| for i in range(num_blocks): |
| self.blocks[i].use_checkpoint = False |
|
|
