Repository: pytorch/PiPPy

Issue:
Pickling error when calling lr_scheduler.step()
```
-- Process 0 terminated with the following error:
Traceback (most recent call last):
  File "/fsx/users/kw2501/conda/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap
    fn(i, *args)
  File "/data/home/kw2501/repos/PiPPy/pippy/utils.py", line 107, in run_worker
    run_master(pp_ranks_per_dp_group[rank], args, *extra_args)
  File "/data/home/kw2501/repos/PiPPy/examples/hf/translation/run_translation.py", line 636, in run_master
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/fsx/users/kw2501/conda/lib/python3.9/site-packages/transformers/trainer.py", line 1505, in train
    return inner_training_loop(
  File "/fsx/users/kw2501/conda/lib/python3.9/site-packages/transformers/trainer.py", line 1817, in _inner_training_loop
    self.lr_scheduler.step()
  File "/data/home/kw2501/repos/PiPPy/pippy/PipelineDriver.py", line 835, in step
    futs.append(scheduler.rpc_async().step(*args, **kwargs))
  File "/fsx/users/kw2501/conda/lib/python3.9/site-packages/torch/distributed/rpc/rref_proxy.py", line 38, in _invoke_rpc
    rref_fut = rref._get_type(timeout=timeout, blocking=False)
  File "/fsx/users/kw2501/conda/lib/python3.9/site-packages/torch/distributed/rpc/api.py", line 429, in _rref_typeof_on_owner
    rref_type = type(rref.local_value())
RuntimeError: Cannot extract tensors from value: AttributeError: Can't pickle local object 'get_linear_schedule_with_warmup.<locals>.lr_lambda'

At:
  /fsx/users/kw2501/conda/lib/python3.9/site-packages/torch/_jit_internal.py(1293): _extract_tensors
```

Repro:
Run on GPU
```
python run_translation.py --model_name_or_path t5-small --do_train --source_lang en --target_lang ro --source_prefix "translate English to Romanian: " --dataset_name wmt16 --dataset_config_name ro-en --output_dir /tmp/tst-translation-gpu --per_device_train_batch_size=8 --per_device_eval_batch_size=8 --overwrite_output_dir --predict_with_generate --max_steps=10  --dp_group_size=1 --pp_group_size=8
```


Discussion:

Resolving patch:
diff --git a/pippy/PipelineDriver.py b/pippy/PipelineDriver.py
--- a/pippy/PipelineDriver.py
+++ b/pippy/PipelineDriver.py
@@ -452,6 +452,8 @@ def __init__(self, stage_id, mod, rank_worker, _record_mem_dumps=False):
         self.optim_init_lock = threading.Lock()
         self.optim_init_cv = threading.Condition(self.optim_init_lock)
 
+        self.lr_scheduler = None
+
     def __getstate__(self):
         # Adding an empty __getstate__ function here to work around the DDP pickling issue (#153) that occurs when the
         # PipelineDiver asks PipeStageExecutors to install_peer_executor(a list of RRefs)
@@ -707,7 +709,11 @@ def instantiate_lr_scheduler(self, lr_sched_class, *args, **kwargs):
                 self.optim_init_cv.wait()
 
         logging.info(f"[{self.stage_id}] Creating learning rate scheduler")
-        return lr_sched_class(self.optimizer, *args, **kwargs)
+        self.lr_scheduler = lr_sched_class(self.optimizer, *args, **kwargs)
+        return self.lr_scheduler
+
+    def step_lr_scheduler(self, *args, **kwargs):
+        self.lr_scheduler.step(*args, **kwargs)
 
     def _check_cleanup(self) -> bool:
         if len(self.value_store):
@@ -822,17 +828,29 @@ def add_param_group(self, param_group):
 
 
 class PipelineLRScheduler(torch.optim.lr_scheduler._LRScheduler):
-    def __init__(self, stage_to_scheds):
+    def __init__(self, stage_to_scheds, stage_to_executor):
         # A dict from stage id to LR schedulers
         self.stage_to_scheds = stage_to_scheds
+        self.stage_to_executor = stage_to_executor
         self.new_step_called = False
         self.last_lr = []
 
     def step(self, *args, **kwargs):
         futs = []
         # Step all remote LR schedulers
+
+        # We use the executor block below because calling scheduler.step()
+        # remotely might cause pickling nested functions, where these nested
+        # functions are usually defined inside user's lr scheduler constructor
+        # as lambda functions to be used by the lr scheduler
+        # See https://github.com/pytorch/PiPPy/issues/404
+        """
         for scheduler in self.stage_to_scheds.values():
             futs.append(scheduler.rpc_async().step(*args, **kwargs))
+        """
+        for executor in self.stage_to_executor.values():
+            futs.append(executor.rpc_async().step_lr_scheduler(*args, **kwargs))
+
         _wait_for_all(futs)
         # Mark new step (invalidates last_lr)
         self.new_step_called = True
@@ -1053,7 +1071,7 @@ def instantiate_lr_scheduler(self, lr_sched_class, *args, **kwargs):
                 remote_lr_sched = executor.remote().instantiate_lr_scheduler(lr_sched_class, *args, **kwargs)
                 stage_to_scheds[stage] = remote_lr_sched
 
-        return PipelineLRScheduler(stage_to_scheds)
+        return PipelineLRScheduler(stage_to_scheds, self.stage_to_executor)
 
     def _sync_replicated_params(self):
         logging.info(f'[root] Synchronizing gradients for {len(self.pipe.replicated_params)} sets of replicated parameters')


Test patch: