remove selective activation checkpointing (#92)

Co-authored-by: Srini Iyer <sviyer@meta.com>
2025-04-22 01:29:09 +00:00 · 2025-04-07 19:20:41 -07:00 · 2025-04-07 19:20:41 -07:00 · 8c1b1a78bb
commit 8c1b1a78bb
parent 1e78a49bf0
1 changed files with 0 additions and 4 deletions
--- a/bytelatent/distributed.py
+++ b/bytelatent/distributed.py
@ -531,10 +531,6 @@ def parallelize_model(
            for i in range(len(module.layers)):
                module.layers[i] = checkpoint_wrapper(
                    module.layers[i],
-                    context_fn=partial(
-                        create_selective_checkpoint_contexts,
-                        get_default_policy(no_recompute_ops),
-                    ),
                )

    if distributed_args.compile: