From 76fbba72b2138dba8f3f09146a3e93e079cfa0f8 Mon Sep 17 00:00:00 2001 From: Srini Iyer Date: Tue, 8 Apr 2025 00:32:48 +0000 Subject: [PATCH] remove selective activation checkpointing --- bytelatent/distributed.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/bytelatent/distributed.py b/bytelatent/distributed.py index 7c99380..d6dc5a5 100644 --- 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: