mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-09 04:50:17 +00:00
[misc] update pre-commit and run all files (#4752)
* [misc] update pre-commit * [misc] run pre-commit * [misc] remove useless configuration files * [misc] ignore cuda for clang-format
This commit is contained in:
@@ -11,13 +11,12 @@ from colossalai.testing import clear_cache_before_run, parameterize
|
||||
from tests.test_auto_parallel.test_offload.model_utils import *
|
||||
|
||||
|
||||
@pytest.mark.skipif(NOT_NVML, reason='pynvml is not installed')
|
||||
@pytest.mark.skipif(NOT_NVML, reason="pynvml is not installed")
|
||||
@clear_cache_before_run()
|
||||
@parameterize('model_name', ['gpt2_', 'bert_'])
|
||||
@parameterize('memory_budget', [4000])
|
||||
@parameterize('solver_name', ['syn', 'asyn'])
|
||||
@parameterize("model_name", ["gpt2_", "bert_"])
|
||||
@parameterize("memory_budget", [4000])
|
||||
@parameterize("solver_name", ["syn", "asyn"])
|
||||
def solver_test(model_name: str, memory_budget: float, solver_name: str):
|
||||
|
||||
get_components_func = non_distributed_component_funcs.get_callable(model_name)
|
||||
model_builder, data_gen = get_components_func()
|
||||
data_args = data_gen(device="cpu")
|
||||
@@ -53,15 +52,15 @@ def solver_test(model_name: str, memory_budget: float, solver_name: str):
|
||||
need_offload = region.need_offload
|
||||
to_prefetch = region.fwd_prefetch_region.r_id if region.fwd_prefetch_region is not None else None
|
||||
print(
|
||||
f'| {model_name} forward | region id: {region.r_id} | need_offload: {need_offload} | to_prefetch: {to_prefetch}'
|
||||
f"| {model_name} forward | region id: {region.r_id} | need_offload: {need_offload} | to_prefetch: {to_prefetch}"
|
||||
)
|
||||
for region in region_list.__reversed__():
|
||||
need_offload = region.need_offload
|
||||
to_prefetch = region.bwd_prefetch_region.r_id if region.bwd_prefetch_region is not None else None
|
||||
print(
|
||||
f'| {model_name} backward | region id: {region.r_id} | need_offload: {need_offload} | to_prefetch: {to_prefetch}'
|
||||
f"| {model_name} backward | region id: {region.r_id} | need_offload: {need_offload} | to_prefetch: {to_prefetch}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
if __name__ == "__main__":
|
||||
solver_test()
|
||||
|
Reference in New Issue
Block a user