MCPcopy Create free account
hub / github.com/deepspeedai/DeepSpeed / _refine_include_states

Method _refine_include_states

deepspeed/runtime/engine.py:5561–5571  ·  view source on GitHub ↗
(self, include: Container[OffloadStateTypeEnum])

Source from the content-addressed store, hash-verified

5559 return self._is_compiled
5560
5561 def _refine_include_states(self, include: Container[OffloadStateTypeEnum]) -> Container[OffloadStateTypeEnum]:
5562 if include is None:
5563 include = list(OffloadStateTypeEnum)
5564
5565 if self.zero_use_cpu_optimizer():
5566 exclude_states = [OffloadStateTypeEnum.hp_params, OffloadStateTypeEnum.optim_states]
5567 if self.zero_optimization_partition_weights():
5568 exclude_states.append(OffloadStateTypeEnum.lp_grads)
5569 include = [x for x in include if x not in exclude_states]
5570
5571 return include
5572
5573 def offload_states(self,
5574 include: Container[OffloadStateTypeEnum] = None,

Callers 1

offload_statesMethod · 0.95

Calls 3

appendMethod · 0.80

Tested by

no test coverage detected