Spaces:
Running
on
Zero
Running
on
Zero
remove selective activation checkpointing (#92)
Browse filesCo-authored-by: Srini Iyer <sviyer@meta.com>
bytelatent/distributed.py
CHANGED
|
@@ -531,10 +531,6 @@ def parallelize_model(
|
|
| 531 |
for i in range(len(module.layers)):
|
| 532 |
module.layers[i] = checkpoint_wrapper(
|
| 533 |
module.layers[i],
|
| 534 |
-
context_fn=partial(
|
| 535 |
-
create_selective_checkpoint_contexts,
|
| 536 |
-
get_default_policy(no_recompute_ops),
|
| 537 |
-
),
|
| 538 |
)
|
| 539 |
|
| 540 |
if distributed_args.compile:
|
|
|
|
| 531 |
for i in range(len(module.layers)):
|
| 532 |
module.layers[i] = checkpoint_wrapper(
|
| 533 |
module.layers[i],
|
|
|
|
|
|
|
|
|
|
|
|
|
| 534 |
)
|
| 535 |
|
| 536 |
if distributed_args.compile:
|