| /external/pytorch/torch/distributed/fsdp/ |
| D | _unshard_param_utils.py | 129 offload_to_cpu: bool,
132 if with_grads and (offload_to_cpu or not state._use_orig_params):
139 if offload_to_cpu and state._handle and (not state._handle.uses_sharded_strategy):
149 if offload_to_cpu and not rank0_only:
164 offload_to_cpu: bool,
172 state, writeback, rank0_only, offload_to_cpu, with_grads
215 if offload_to_cpu and handle.uses_sharded_strategy:
241 offload_to_cpu: bool,
245 state, writeback, rank0_only, offload_to_cpu, with_grads
261 offload_to_cpu=offload_to_cpu,
[all …]
|
| D | api.py | 288 offload_to_cpu: bool = False variable in StateDictConfig
367 offload_to_cpu: bool = True variable in OptimStateDictConfig
384 offload_to_cpu: bool = False variable in LocalOptimStateDictConfig
|
| D | _state_dict_utils.py | 114 offload_to_cpu: bool = False,
131 offload_to_cpu=offload_to_cpu,
160 offload_to_cpu: bool,
174 offload_to_cpu=offload_to_cpu,
258 fsdp_state._state_dict_config.offload_to_cpu
306 offload_to_cpu=fsdp_state._state_dict_config.offload_to_cpu,
446 if fsdp_state._state_dict_config.offload_to_cpu:
530 offload_to_cpu=False,
565 if fsdp_state._state_dict_config.offload_to_cpu:
|
| D | fully_sharded_data_parallel.py | 611 offload_to_cpu=False,
876 offload_to_cpu: bool = False,
940 module, recurse, writeback, rank0_only, offload_to_cpu, with_grads
|
| /external/pytorch/test/distributed/fsdp/ |
| D | test_fsdp_dtensor_state_dict.py | 115 self, offload_to_cpu, is_even_sharded_model argument
123 state_dict_config=ShardedStateDictConfig(offload_to_cpu=offload_to_cpu),
125 offload_to_cpu=offload_to_cpu
135 state_dict_config=ShardedStateDictConfig(offload_to_cpu=offload_to_cpu),
137 offload_to_cpu=offload_to_cpu
197 self, offload_to_cpu, is_even_sharded_model argument
206 offload_to_cpu=offload_to_cpu
255 self, offload_to_cpu, is_even_sharded_model argument
263 state_dict_config=ShardedStateDictConfig(offload_to_cpu=offload_to_cpu),
|
| D | test_hsdp_dtensor_state_dict.py | 110 def test_dtensor_sharded_tensor_state_dict_identical(self, offload_to_cpu): argument
117 state_dict_config=ShardedStateDictConfig(offload_to_cpu=offload_to_cpu),
119 offload_to_cpu=offload_to_cpu
129 state_dict_config=ShardedStateDictConfig(offload_to_cpu=offload_to_cpu),
131 offload_to_cpu=offload_to_cpu
179 def test_dtensor_sharded_optim_load_state_dict(self, offload_to_cpu): argument
187 offload_to_cpu=offload_to_cpu
233 def test_dtensor_sharded_model_load_state_dict(self, offload_to_cpu): argument
240 state_dict_config=ShardedStateDictConfig(offload_to_cpu=offload_to_cpu),
|
| D | test_fsdp_unshard_params.py | 120 offload_to_cpu: bool,
158 offload_to_cpu=offload_to_cpu,
166 offload_to_cpu
283 offload_to_cpu: bool,
329 offload_to_cpu=offload_to_cpu,
348 offload_to_cpu=offload_to_cpu,
486 offload_to_cpu: bool,
505 offload_to_cpu=offload_to_cpu,
525 sharding_strategy == ShardingStrategy.NO_SHARD and not offload_to_cpu
553 is_supported = use_orig_params and not offload_to_cpu
|
| D | test_fsdp_state_dict.py | 279 offload_to_cpu=state_dict_rank0_and_offload,
283 offload_to_cpu=state_dict_rank0_and_offload,
287 offload_to_cpu=state_dict_rank0_and_offload,
334 offload_to_cpu=False,
|
| D | test_fsdp_optim_state.py | 1883 ShardedOptimStateDictConfig(offload_to_cpu=False),
1900 offload_to_cpu=True,
|
| /external/pytorch/torch/testing/_internal/distributed/ |
| D | common_state_dict.py | 23 def _compare_tensor(self, orig_tensor, dist_tensor, offload_to_cpu=False): argument
27 if offload_to_cpu:
38 offload_to_cpu=False, argument
47 self._compare_tensor(param, dist_param, offload_to_cpu)
|
| /external/pytorch/torch/distributed/_composable/fsdp/ |
| D | _fsdp_param.py | 226 self.offload_to_cpu: bool = isinstance(offload_policy, CPUOffloadPolicy)
228 self.offload_to_cpu and cast(CPUOffloadPolicy, offload_policy).pin_memory
344 if self.offload_to_cpu and not padded_sharded_param.is_meta:
624 if self.offload_to_cpu:
637 if self.offload_to_cpu:
|
| D | _fsdp_collectives.py | 193 and not fsdp_param.offload_to_cpu
379 if fsdp_param.offload_to_cpu:
|
| /external/pytorch/benchmarks/dynamo/ |
| D | dist_util.py | 117 offload_to_cpu=False,
|
| /external/pytorch/test/distributed/checkpoint/e2e/ |
| D | test_e2e_save_and_load.py | 370 self._verify_msd(model_sd, loaded_model_sd, offload_to_cpu=True)
379 loaded_optim_state[k][optim_key], v[optim_key], offload_to_cpu=True
|
| /external/pytorch/torch/distributed/checkpoint/ |
| D | state_dict.py | 337 offload_to_cpu=options.cpu_offload, rank0_only=options.cpu_offload
340 offload_to_cpu=options.cpu_offload,
346 offload_to_cpu=options.cpu_offload,
349 offload_to_cpu=options.cpu_offload,
|