| 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
|
|
|
