Repository: microsoft/DeepSpeed

Issue:
DeepSpeed v1.6.1 issues with Pytorch Lightning Deep Speed 3 Offloading
**Describe the bug**
DeepSpeed Stage 3 with Offloading is throwing an error. This error is not thrown in DeepSpeed version 0.5.8.  (See Screenshots section). I am currently using PyTorch 1.8.2, but the error is also thrown in PyTorch 1.11.0.

**To Reproduce**
I am using a BERT model from transformers v.4.17 with a Sequence Classification tuning head. I have wrapped this model in a Pytorch-Lightning module. I am training this model with the `deepspeed_stage_3_offload` strategy. 

```
from pytorch_lightning import Trainer

model = get_model()
dm = get_data_module()

trainer = Trainer(
        max_epochs=10,
        accelerator="gpu",
        devices=3,
        strategy="deepspeed_stage_3_offload",
        auto_lr_find=True,
        logger=logger,
)
trainer.fit(model, dm)
```

**Expected behavior**
The model trains 

**ds_report output**
```
❯ ds_report
[2022-04-02 23:14:34,784] [WARNING] [partition_parameters.py:54:<module>] unable to find torch.distributed._all_gather_base. will fall back to torch.distributed.all_gather which will result in suboptimal performance. please consider upgrading your pytorch installation.
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
 [WARNING]  using untested triton version (1.1.1), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
async_io ............... [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/home/keshav/mambaforge/envs/old_pytorch/lib/python3.7/site-packages/torch']
torch version .................... 1.8.2+cu111
torch cuda version ............... 11.1
torch hip version ................ None
nvcc version ..................... 11.1
deepspeed install path ........... ['/home/keshav/mambaforge/envs/old_pytorch/lib/python3.7/site-packages/deepspeed']
deepspeed info ................... 0.6.1, unknown, unknown
deepspeed wheel compiled w. ...... torch 1.8, cuda 11.1, hip 0.0
```

**Screenshots**
```
Traceback (most recent call last):
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/runpy.py", line 197, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/mnt/storage/grid/home/keshav/projects/Ibis/experiments/dependency_tests/training_with_deepspeed_stage_3_offload.py", line 48, in <module>
    train()
  File "/mnt/storage/grid/home/keshav/projects/Ibis/experiments/dependency_tests/training_with_deepspeed_stage_3_offload.py", line 39, in train
    trainer.fit(model, dm)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 771, in fit
    self._call_and_handle_interrupt(
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 722, in _call_and_handle_interrupt
    return self.strategy.launcher.launch(trainer_fn, *args, trainer=self, **kwargs)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 93, in launch
    return function(*args, **kwargs)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 812, in _fit_impl
    results = self._run(model, ckpt_path=self.ckpt_path)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1237, in _run
    results = self._run_stage()
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1324, in _run_stage
    return self._run_train()
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1354, in _run_train
    self.fit_loop.run()
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/loops/fit_loop.py", line 269, in advance
    self._outputs = self.epoch_loop.run(self._data_fetcher)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py", line 208, in advance
    batch_output = self.batch_loop.run(batch, batch_idx)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/loops/batch/training_batch_loop.py", line 88, in advance
    outputs = self.optimizer_loop.run(split_batch, optimizers, batch_idx)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 203, in advance
    result = self._run_optimization(
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 256, in _run_optimization
    self._optimizer_step(optimizer, opt_idx, batch_idx, closure)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 369, in _optimizer_step
    self.trainer._call_lightning_module_hook(
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1596, in _call_lightning_module_hook
    output = fn(*args, **kwargs)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/core/lightning.py", line 1625, in optimizer_step
    optimizer.step(closure=optimizer_closure)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/core/optimizer.py", line 168, in step
    step_output = self._strategy.optimizer_step(self._optimizer, self._optimizer_idx, closure, **kwargs)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/strategies/ddp.py", line 278, in optimizer_step
    optimizer_output = super().optimizer_step(optimizer, opt_idx, closure, model, **kwargs)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/strategies/strategy.py", line 193, in optimizer_step
    return self.precision_plugin.optimizer_step(model, optimizer, opt_idx, closure, **kwargs)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/plugins/precision/deepspeed.py", line 70, in optimizer_step
    closure_result = closure()
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 148, in __call__
    self._result = self.closure(*args, **kwargs)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 143, in closure
    self._backward_fn(step_output.closure_loss)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 311, in backward_fn
    self.trainer._call_strategy_hook("backward", loss, optimizer, opt_idx)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py", line 1766, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/strategies/strategy.py", line 168, in backward
    self.precision_plugin.backward(self.lightning_module, closure_loss, *args, **kwargs)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/pytorch_lightning/plugins/precision/deepspeed.py", line 51, in backward
    deepspeed_engine.backward(closure_loss, *args, **kwargs)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/deepspeed/utils/nvtx.py", line 11, in wrapped_fn
    return func(*args, **kwargs)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/deepspeed/runtime/engine.py", line 1667, in backward
    self.optimizer.backward(loss)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/deepspeed/utils/nvtx.py", line 11, in wrapped_fn
    return func(*args, **kwargs)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/deepspeed/runtime/zero/stage3.py", line 2793, in backward
    self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/deepspeed/runtime/fp16/loss_scaler.py", line 53, in backward
    scaled_loss.backward(retain_graph=retain_graph)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/torch/_tensor.py", line 363, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/torch/autograd/__init__.py", line 173, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/torch/autograd/function.py", line 253, in apply
    return user_fn(self, *args)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/deepspeed/runtime/zero/stage3.py", line 510, in backward
    ctx.pre_backward_function(ctx.module)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/deepspeed/utils/nvtx.py", line 11, in wrapped_fn
    return func(*args, **kwargs)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/deepspeed/runtime/zero/stage3.py", line 1461, in _run_before_backward_function
    self.pre_sub_module_backward_function(sub_module)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/deepspeed/runtime/zero/stage3.py", line 1549, in pre_sub_module_backward_function
    self.param_coordinator.fetch_sub_module(sub_module)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/deepspeed/utils/nvtx.py", line 11, in wrapped_fn
    return func(*args, **kwargs)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/keshav/mambaforge/envs/dl/lib/python3.9/site-packages/deepspeed/runtime/zero/stage3.py", line 358, in fetch_sub_module
    raise RuntimeError(
RuntimeError: tracing error at step 1429: expected the next 2 parameters in the parameter fetch queue to be ({'id': 487, 'status': 'AVAILABLE', 'numel': 1455104, 'ds_numel': 1455104, 'shape': (1421, 1024), 'ds_shape': (1421, 1024), 'requires_grad': True, 'grad_shape': None, 'persist': False, 'active_sub_modules': {557}}, {'id': 488, 'status': 'AVAILABLE', 'numel': 1421, 'ds_numel': 1421, 'shape': (1421,), 'ds_shape': (1421,), 'requires_grad': True, 'grad_shape': None, 'persist': True, 'active_sub_modules': {557}}) but got ().
```

**System info (please complete the following information):**
 - OS: Ubuntu 20.03
 -  GPUs: 8 TU104GL [Quadro RTX 5000]
 - Single Machine
 - Python version: 3.7
 - Huggingface's Transformers 4.17
 - DeepSpeed 0.6.1
 - Pytorch-Lightning 1.6.0
 - PyTorch 1.8.2

**Launcher context**
Are you launching your experiment with the `deepspeed` launcher, MPI, or something else?

I am launching with PyTorch-Lightning. I have written a function that is executed via a python script.

**Docker context**
Are you using a specific docker image that you can share?
n/a

**Additional context**
n/a



Discussion:
@keshavd, to help with repro can you please provide your full training script that includes  `get_model()` and `get_data_module()`?
I can provide an example @tjruwase. I will also add that using the commented out version of `configure_optimizers` in the script below (which I got from a PTL example) results in the following traceback:

```
"expected input to be on device CUDA" error
```

```
from datetime import datetime
from typing import Optional
from argparse import ArgumentParser

import datasets
import torch
from pytorch_lightning import LightningDataModule, LightningModule, Trainer, seed_everything
from torch.utils.data import DataLoader
from transformers import (
    AdamW,
    AutoConfig,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    get_linear_schedule_with_warmup,
)

from deepspeed.ops.adam import FusedAdam

AVAIL_GPUS = min(1, torch.cuda.device_count())

class GLUEDataModule(LightningDataModule):

    task_text_field_map = {
        "cola": ["sentence"],
        "sst2": ["sentence"],
        "mrpc": ["sentence1", "sentence2"],
        "qqp": ["question1", "question2"],
        "stsb": ["sentence1", "sentence2"],
        "mnli": ["premise", "hypothesis"],
        "qnli": ["question", "sentence"],
        "rte": ["sentence1", "sentence2"],
        "wnli": ["sentence1", "sentence2"],
        "ax": ["premise", "hypothesis"],
    }

    glue_task_num_labels = {
        "cola": 2,
        "sst2": 2,
        "mrpc": 2,
        "qqp": 2,
        "stsb": 1,
        "mnli": 3,
        "qnli": 2,
        "rte": 2,
        "wnli": 2,
        "ax": 3,
    }

    loader_columns = [
        "datasets_idx",
        "input_ids",
        "token_type_ids",
        "attention_mask",
        "start_positions",
        "end_positions",
        "labels",
    ]

    def __init__(
            self,
            model_name_or_path: str,
            task_name: str = "mrpc",
            max_seq_length: int = 128,
            train_batch_size: int = 32,
            eval_batch_size: int = 32,
            **kwargs,
    ):
        super().__init__()
        self.model_name_or_path = model_name_or_path
        self.task_name = task_name
        self.max_seq_length = max_seq_length
        self.train_batch_size = train_batch_size
        self.eval_batch_size = eval_batch_size

        self.text_fields = self.task_text_field_map[task_name]
        self.num_labels = self.glue_task_num_labels[task_name]
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path, use_fast=True)

    def setup(self, stage: str):
        self.dataset = datasets.load_dataset("glue", self.task_name)

        for split in self.dataset.keys():
            self.dataset[split] = self.dataset[split].map(
                self.convert_to_features,
                batched=True,
                remove_columns=["label"],
            )
            self.columns = [c for c in self.dataset[split].column_names if c in self.loader_columns]
            self.dataset[split].set_format(type="torch", columns=self.columns)

        self.eval_splits = [x for x in self.dataset.keys() if "validation" in x]

    def prepare_data(self):
        datasets.load_dataset("glue", self.task_name)
        AutoTokenizer.from_pretrained(self.model_name_or_path, use_fast=True)

    def train_dataloader(self):
        return DataLoader(self.dataset["train"], batch_size=self.train_batch_size)

    def val_dataloader(self):
        if len(self.eval_splits) == 1:
            return DataLoader(self.dataset["validation"], batch_size=self.eval_batch_size)
        elif len(self.eval_splits) > 1:
            return [DataLoader(self.dataset[x], batch_size=self.eval_batch_size) for x in self.eval_splits]

    def test_dataloader(self):
        if len(self.eval_splits) == 1:
            return DataLoader(self.dataset["test"], batch_size=self.eval_batch_size)
        elif len(self.eval_splits) > 1:
            return [DataLoader(self.dataset[x], batch_size=self.eval_batch_size) for x in self.eval_splits]

    def convert_to_features(self, example_batch, indices=None):

        # Either encode single sentence or sentence pairs
        if len(self.text_fields) > 1:
            texts_or_text_pairs = list(zip(example_batch[self.text_fields[0]], example_batch[self.text_fields[1]]))
        else:
            texts_or_text_pairs = example_batch[self.text_fields[0]]

        # Tokenize the text/text pairs
        features = self.tokenizer.batch_encode_plus(
            texts_or_text_pairs, max_length=self.max_seq_length, pad_to_max_length=True, truncation=True
        )

        # Rename label to labels to make it easier to pass to model forward
        features["labels"] = example_batch["label"]

        return features


class GLUETransformer(LightningModule):
    def __init__(
            self,
            model_name_or_path: str,
            num_labels: int,
            task_name: str,
            learning_rate: float = 2e-5,
            adam_epsilon: float = 1e-8,
            warmup_steps: int = 0,
            weight_decay: float = 0.0,
            train_batch_size: int = 8,
            eval_batch_size: int = 8,
            eval_splits: Optional[list] = None,
            **kwargs,
    ):
        super().__init__()

        self.save_hyperparameters()

        self.config = AutoConfig.from_pretrained(model_name_or_path, num_labels=num_labels)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_name_or_path, config=self.config)
        self.metric = datasets.load_metric(
            "glue", self.hparams.task_name, experiment_id=datetime.now().strftime("%d-%m-%Y_%H-%M-%S")
        )

    def forward(self, **inputs):
        return self.model(**inputs)

    def training_step(self, batch, batch_idx):
        outputs = self(**batch)
        loss = outputs[0]
        return loss

    def validation_step(self, batch, batch_idx, dataloader_idx=0):
        outputs = self(**batch)
        val_loss, logits = outputs[:2]

        if self.hparams.num_labels >= 1:
            preds = torch.argmax(logits, axis=1)
        elif self.hparams.num_labels == 1:
            preds = logits.squeeze()

        labels = batch["labels"]

        return {"loss": val_loss, "preds": preds, "labels": labels}

    def validation_epoch_end(self, outputs):
        if self.hparams.task_name == "mnli":
            for i, output in enumerate(outputs):
                # matched or mismatched
                split = self.hparams.eval_splits[i].split("_")[-1]
                preds = torch.cat([x["preds"] for x in output]).detach().cpu().numpy()
                labels = torch.cat([x["labels"] for x in output]).detach().cpu().numpy()
                loss = torch.stack([x["loss"] for x in output]).mean()
                self.log(f"val_loss_{split}", loss, prog_bar=True)
                split_metrics = {
                    f"{k}_{split}": v for k, v in self.metric.compute(predictions=preds, references=labels).items()
                }
                self.log_dict(split_metrics, prog_bar=True)
            return loss

        preds = torch.cat([x["preds"] for x in outputs]).detach().cpu().numpy()
        labels = torch.cat([x["labels"] for x in outputs]).detach().cpu().numpy()
        loss = torch.stack([x["loss"] for x in outputs]).mean()
        self.log("val_loss", loss, prog_bar=True)
        self.log_dict(self.metric.compute(predictions=preds, references=labels), prog_bar=True)
        return loss

    def setup(self, stage=None) -> None:
        if stage != "fit":
            return
        # Get dataloader by calling it - train_dataloader() is called after setup() by default
        train_loader = self.trainer.datamodule.train_dataloader()

        # Calculate total steps
        tb_size = self.hparams.train_batch_size * max(1, self.trainer.gpus)
        ab_size = self.trainer.accumulate_grad_batches * float(self.trainer.max_epochs)
        self.total_steps = (len(train_loader.dataset) // tb_size) // ab_size

    def configure_optimizers(self):
        """Prepare optimizer and schedule (linear warmup and decay)"""
        # Stage 3 offload will run with this version of configure_optimizers
        model = self.model
        no_decay = ["bias", "LayerNorm.weight"]
        optimizer_grouped_parameters = [
            {
                "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
                "weight_decay": self.hparams.weight_decay,
            },
            {
                "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
                "weight_decay": 0.0,
            },
        ]
        optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=self.hparams.learning_rate, eps=self.hparams.adam_epsilon)
    
        scheduler = get_linear_schedule_with_warmup(
            optimizer,
            num_warmup_steps=self.hparams.warmup_steps,
            num_training_steps=self.total_steps,
        )
        scheduler = {"scheduler": scheduler, "interval": "step", "frequency": 1}
        return [optimizer], [scheduler]

    # NOTE: this results in the "expected input to be on device CUDA" error
    # def configure_optimizers(self):
    #     return FusedAdam(self.parameters())


    seed_everything(42)


# https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html
def main(args):

    print(args)
    dm = GLUEDataModule(model_name_or_path=args.model_architecture,
                        task_name="cola",
                        train_batch_size=args.batch_size,
                        eval_batch_size=args.batch_size)
    dm.setup("fit")
    model = GLUETransformer(
        model_name_or_path=args.model_architecture,
        num_labels=dm.num_labels,
        eval_splits=dm.eval_splits,
        task_name=dm.task_name,
    )

    trainer = Trainer.from_argparse_args(args)
    trainer.fit(model, datamodule=dm)
    # trainer.test()
    # trainer.predict()


if __name__ == "__main__":

    # run this from the command line using this template:
    # TOKENIZERS_PARALLELISM="true" python pytorch_lightning/run_ptl.py --accelerator=gpu --devices=4 --strategy="deepspeed_stage_3" --precision=32 --batch_size=1 --max_epochs=1 --model_architecture="allenai/longformer-base-4096"

    parser = ArgumentParser()
    parser.add_argument("--accelerator", default="gpu")
    parser.add_argument("--devices", default=4, type=int)
    parser.add_argument("--strategy", default=None)
    parser.add_argument("--batch_size", default=1, type=int)
    # options: '16', '32', '64', 'bf16', 'mixed'
    parser.add_argument("--precision", default=32)
    parser.add_argument("--max_epochs", default=1, type=int)
    parser.add_argument("--model_architecture", default=None)

    # parser = Trainer.add_argparse_args(parser)


    args = parser.parse_args()
    # args = parser.parse_known_args()
    main(args)
```
    
I run it with this command:

```
TOKENIZERS_PARALLELISM="true" python pytorch_lightning/run_ptl.py --accelerator=gpu --devices=4 --strategy="deepspeed_stage_3_offload" --precision=16 --batch_size=16 --max_epochs=1 --model_architecture="allenai/longformer-base-4096"````

Just tested with `deepspeed==0.5.4` and the training completes successfully. So this was introduced somewhere between `0.5.4` and `0.6.1`
here are compat results with the versions


:heavy_check_mark: `0.5.4`
:heavy_check_mark: `0.5.5`
:heavy_check_mark: `0.5.6`
:heavy_check_mark: `0.5.7`
:heavy_check_mark: `0.5.8`
:heavy_check_mark: `0.5.9`
:heavy_check_mark: `0.5.10`
:heavy_multiplication_x: `0.6.0`
:heavy_multiplication_x: `0.6.1`

`0.6.x` versions all result in this traceback for me:

```
RuntimeError: tracing error at step 390: expected the next 2 parameters in the parameter fetch queue to be ({'id': 271, 'status': 'AVAILABLE', 'numel': 1536, 'ds_numel': 1536, 'shape': (2, 768), 'ds_shape': (2, 768), 'requires_grad': True, 'grad_shape': None, 'persist': True, 'active_sub_modules': {255}}, {'id': 272, 'status': 'AVAILABLE', 'numel': 2, 'ds_numel': 2, 'shape': (2,), 'ds_shape': (2,), 'requires_grad': True, 'grad_shape': None, 'persist': True, 'active_sub_modules': {255}}) but got ().
```
@tjruwase 
Thanks @jmwoloso for opening this issue and providing a reproduce! I'll try to help out here as well @tjruwase.
@jmwoloso, thanks for the repro guide. I am trying it on my side. Thanks, @SeanNaren. 
@jmwoloso, @SeanNaren, any idea how to get past this?
```bash
Downloading and preparing dataset glue/cola (download: 368.14 KiB, generated: 596.73 KiB, post-processed: Unknown size, total: 964.86 KiB) to /home/olruwase/.cache/huggingface/datasets/glue/cola/1.0.0/dacbe3125aa31d7f70367a07a8a9e72a5a0bfeb5fc42e75c9db75b96da6053ad...
Traceback (most recent call last):
  File "run.py", line 285, in <module>
    main(args)
  File "run.py", line 251, in main
    dm.setup("fit")
  File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/core/datamodule.py", line 474, in wrapped_fn
    fn(*args, **kwargs)
  File "run.py", line 80, in setup
    self.dataset = datasets.load_dataset("glue", self.task_name)
  File "/home/olruwase/.local/lib/python3.6/site-packages/datasets/load.py", line 1117, in load_dataset
    use_auth_token=use_auth_token,
  File "/home/olruwase/.local/lib/python3.6/site-packages/datasets/builder.py", line 637, in download_and_prepare
    dl_manager=dl_manager, verify_infos=verify_infos, **download_and_prepare_kwargs
  File "/home/olruwase/.local/lib/python3.6/site-packages/datasets/builder.py", line 704, in _download_and_prepare
    split_generators = self._split_generators(dl_manager, **split_generators_kwargs)
  File "/home/olruwase/.cache/huggingface/modules/datasets_modules/datasets/glue/dacbe3125aa31d7f70367a07a8a9e72a5a0bfeb5fc42e75c9db75b96da6053ad/glue.py", line 490, in _split_generators
    dl_dir = dl_manager.download_and_extract(self.config.data_url)
  File "/home/olruwase/.local/lib/python3.6/site-packages/datasets/utils/download_manager.py", line 284, in download_and_extract
    return self.extract(self.download(url_or_urls))
  File "/home/olruwase/.local/lib/python3.6/site-packages/datasets/utils/download_manager.py", line 197, in download
    download_func, url_or_urls, map_tuple=True, num_proc=download_config.num_proc, disable_tqdm=False
  File "/home/olruwase/.local/lib/python3.6/site-packages/datasets/utils/py_utils.py", line 197, in map_nested
    return function(data_struct)
  File "/home/olruwase/.local/lib/python3.6/site-packages/datasets/utils/download_manager.py", line 217, in _download
    return cached_path(url_or_filename, download_config=download_config)
  File "/home/olruwase/.local/lib/python3.6/site-packages/datasets/utils/file_utils.py", line 299, in cached_path
    use_auth_token=download_config.use_auth_token,
  File "/home/olruwase/.local/lib/python3.6/site-packages/datasets/utils/file_utils.py", line 606, in get_from_cache
    raise ConnectionError("Couldn't reach {}".format(url))
ConnectionError: Couldn't reach https://dl.fbaipublicfiles.com/glue/data/CoLA.zip
```
> ```shell
>   File "/home/olruwase/.local/lib/python3.6/site-packages/datasets/utils/file_utils.py", line 606, in get_from_cache
>     raise ConnectionError("Couldn't reach {}".format(url))
> ConnectionError: Couldn't reach https://dl.fbaipublicfiles.com/glue/data/CoLA.zip
> ```

Never mind, connection error resolved itself. 
Okay, I got a repro
```bash
  File "/opt/conda/lib/python3.6/site-packages/deepspeed/runtime/zero/stage3.py", line 359, in fetch_sub_module
    f"tracing error at step {self.__step_id}: "
RuntimeError: tracing error at step 246: expected the next 2 parameters in the parameter fetch queue to be ({'id': 271, 'status': 'AVAILABLE', 'numel': 1536, 'ds_numel': 1536, 'shape': (2, 768), 'ds_shape': (2, 768), 'requires_grad': True, 'grad_shape': None, 'persist': True, 'active_sub_modules': {243}}, {'id': 272, 'status': 'AVAILABLE', 'numel': 2, 'ds_numel': 2, 'shape': (2,), 'ds_shape': (2,), 'requires_grad': True, 'grad_shape': None, 'persist': True, 'active_sub_modules': {243}}) but got ().
```

Resolving patch:
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -1568,10 +1568,6 @@ def forward(self, *inputs, **kwargs):
         loss = self.module(*inputs, **kwargs)
 
         if self.zero_optimization_partition_weights():
-            # Reset the ZeRO-3 state if we are only doing forward-passes (ie evaluation).
-            if not torch._C.is_grad_enabled():
-                self.optimizer.param_coordinator.reset_step()
-
             # Disable automated discovery of external parameters
             for module in self.module.modules():
                 module._parameters._in_forward = False
diff --git a/deepspeed/runtime/zero/partitioned_param_coordinator.py b/deepspeed/runtime/zero/partitioned_param_coordinator.py
new file mode 100644
--- /dev/null
+++ b/deepspeed/runtime/zero/partitioned_param_coordinator.py
@@ -0,0 +1,357 @@
+"""
+"Copyright 2020 The Microsoft DeepSpeed Team.
+Licensed under the MIT license.
+"""
+
+from dataclasses import dataclass
+import functools
+import collections
+from collections import OrderedDict, UserDict
+from typing import Deque, Dict, Iterable, Set, Tuple
+import torch
+from torch.cuda import Event, Stream
+from torch.nn import Module, Parameter
+
+from deepspeed.utils.logging import logger
+from deepspeed.runtime.zero.partition_parameters import *
+from deepspeed.runtime.zero.offload_constants import *
+from deepspeed.runtime.swap_tensor.partitioned_param_swapper import PartitionedParamStatus
+
+
+def debug_rank0(message: str) -> None:
+    if dist.get_rank() == 0:
+        logger.debug(message)
+
+
+@instrument_w_nvtx
+def get_all_parameters(sub_module, recurse=False):
+    return itertools.chain(sub_module.named_parameters(recurse=recurse),
+                           sub_module.ds_external_parameters())
+
+
+def iter_params(module: Module, recurse=False) -> Iterable[Parameter]:
+    return map(lambda pair: pair[1], get_all_parameters(module, recurse))
+
+
+class PartitionedParameterCoordinator:
+    """Handles partitioning and gathering of parameters."""
+    class __InflightParamRegistry(UserDict):
+        """registry for parameters in flight"""
+        def __setitem__(self,
+                        param: Parameter,
+                        handle: AllGatherCoalescedHandle) -> None:
+            if param in self.data:
+                raise RuntimeError(f"{param.ds_summary()} already in registry")
+            if param.ds_status != ZeroParamStatus.INFLIGHT:
+                raise RuntimeError(
+                    f"attempted to add non-inflight parameter to registry {param.ds_summary()}"
+                )
+            self.data[param] = handle
+
+    @dataclass
+    class __ParamInTrace:
+        param: Parameter
+        step_id_last_used_at: int
+
+    def __init__(
+        self,
+        prefetch_bucket_sz: int,
+        max_reuse_distance_in_numel: int,
+        max_available_parameters_in_numel: int,
+        allgather_stream: Stream,
+        prefetch_nvme: bool = False,
+    ) -> None:
+        # mapping of param -> handle for each param that is currently in flight
+        self.__inflight_param_registry = __class__.__InflightParamRegistry()
+        # keeps track of the number of submodules invoked so far.
+        self.__step_id: int = 0
+        # whether or not we have completed a trace of the entire network. This should
+        # always be true after the first forward pass + backward pass.
+        self.trace_complete: bool = False
+        # sequence of submodules/parameters in forward pass + backward pass
+        self.__submodule_order: Iterable[Module] = []
+        self.__param_order: Iterable[__class__.__ParamInTrace] = []
+        self.__most_recent_step_id_param_fetched_for = collections.defaultdict(
+            lambda: int(-1e10))
+        # number of available params, and max number of available params
+        self.__n_available_params: int = 0
+        self.__max_n_available_params: int = max_available_parameters_in_numel
+        # max distance between two use of the module beyond which module is released
+        self.__max_reuse_dist_in_numel: int = max_reuse_distance_in_numel
+        # queue for parameters to fetch. parameters will be popped off the left
+        # side of the dequeue as they are fetched
+        self.__param_queue: Deque[__class__.__ParamInTrace] = None
+        self.__prefetch_bucket_sz: int = prefetch_bucket_sz
+        self.__prefetch_nvme: bool = prefetch_nvme
+        self.hierarchy: int = 0
+
+        # stream that will be used for allgather operations
+        self.__allgather_stream: Stream = allgather_stream
+
+        # limit the number of fetch events that can be queued at once
+        # otherwise, what happens is memory is allocated by the host thread at the
+        # time of the call, but not used until later by the asynchronous cuda stream.
+        # allowing an infinite number of these to queue up causes a lot of memory
+        # pressure that then becomes detrimental to performance.
+        # this is a much less elegant way of fixing this vs something like using
+        # cudaMallocAsync/cudaFreeAsync. Choosing to not expose this to the user now
+        # because ideally in the future its replaced by an async allocation
+        # mechanism which doesn't require any configuration by the user.
+        self.__ongoing_fetch_events: Deque[Event] = collections.deque()
+        # TODO. make this configurable via JSON
+        self.__max_ongoing_fetch_events: int = 2
+
+    """Tracing and Tracking
+    TODO. consider performing trace before initializing PartitionedParameterCoordinator
+    and passing trace results into constructor. This way all the code in here can
+    just assume that the trace is complete and the results can be entirely
+    immutable.
+
+    Bookkeeping operations used to track where we are in the forward/backward pass
+    """
+
+    def record_trace(self, sub_module: Module) -> None:
+        """adds sub module to trace"""
+        if self.trace_complete:
+            raise RuntimeError(
+                "attempted to record trace when trace was already complete")
+
+        self.__submodule_order.append(sub_module)
+        for param in sorted(set(iter_params(sub_module)), key=lambda p: p.ds_id):
+            self.__param_order.append(
+                __class__.__ParamInTrace(param=param,
+                                         step_id_last_used_at=self.__step_id))
+
+    def reset_step(self) -> None:
+        """indicate that we have completed one fwd+bwd for the model"""
+        if self.__inflight_param_registry:
+            raise RuntimeError(
+                f"still have inflight params "
+                f"{[p.ds_summary for p in self.__inflight_param_registry.keys()]}")
+
+        if not self.trace_complete:
+            # make sure that recorded parameter and submodule orders are
+            # identical across ranks
+            assert_ints_same_as_other_ranks([m.id for m in self.__submodule_order])
+            assert_ints_same_as_other_ranks([p.param.ds_id for p in self.__param_order])
+            assert_ints_same_as_other_ranks(
+                [p.step_id_last_used_at for p in self.__param_order])
+
+            self.__submodule_order = tuple(self.__submodule_order)  # freeze
+            self.__param_order = tuple(self.__param_order)  # freeze
+            self.trace_complete = True
+            print_rank_0(f"completed trace: {[m.id for m in self.__submodule_order]}",
+                         force=False)
+
+        self.__param_queue = collections.deque(self.__param_order)  # reset fetch queue
+        self.__most_recent_step_id_param_fetched_for = collections.defaultdict(
+            lambda: int(-1e10))
+        self.__step_id = 0
+        self.__n_available_params = 0
+
+    """Fetch and Release
+    Fetching, prefetching, and releasing parameters
+    """
+
+    @instrument_w_nvtx
+    @torch.no_grad()
+    def fetch_sub_module(self, current_submodule: Module) -> None:
+        """This method does the following (in order):
+        1. kick off fetch for parameters in immediately required sub module
+        2. kick off fetch for next few parameters we will need later (prefetch)
+        3. block on parameters in immediately required sub module
+        """
+        debug_rank0(
+            f"{self.__step_id}: M{current_submodule.id}({type(current_submodule).__name__}) P{[p.ds_id for p in iter_params(current_submodule)]} "
+            + str({
+                "avail": f"{self.__n_available_params:.1e}",
+                "queue_sz": f"{len(self.__param_queue or [])}",
+                "inflight": [p.ds_id for p in self.__inflight_param_registry],
+            }))
+
+        params_to_fetch = frozenset(iter_params(current_submodule))
+
+        # kick off all gather for params in the immediately required submodule
+        for param in params_to_fetch:
+            debug_rank0(f"-fetch: {param.ds_summary()}")
+        self.__all_gather_params(params_to_fetch)
+
+        # wait for parameters in the immediately needed submodule to become available
+        for param in params_to_fetch:
+            param.ds_active_sub_modules.add(current_submodule.id)
+            debug_rank0(f"-wait: {param.ds_summary()}")
+            if param in self.__inflight_param_registry:
+                with torch.cuda.stream(self.__allgather_stream):
+                    while self.__ongoing_fetch_events and self.__ongoing_fetch_events[
+                            0].query():
+                        self.__ongoing_fetch_events.popleft()
+                    if len(self.__ongoing_fetch_events
+                           ) > self.__max_ongoing_fetch_events:
+                        self.__ongoing_fetch_events.popleft().synchronize()
+
+                    self.__inflight_param_registry.pop(param).wait()
+
+                    event = Event()
+                    event.record()
+                    self.__ongoing_fetch_events.append(event)
+
+            assert param.ds_status == ZeroParamStatus.AVAILABLE, param.ds_summary()
+        torch.cuda.current_stream().wait_stream(self.__allgather_stream)
+
+        # kick off parameter prefetches for upcoming modules
+        # don't prefetch if we dont have a completed model trace, or if we aren't
+        # training (throws off the tracing and don't want to prefetch modules for bwd)
+        if self.trace_complete and current_submodule.training:
+            # go through the parameters we need for the current module and pop them
+            # off the fetch queue so that they aren't prefetched later.
+            # if params have already been popped off the fetch queue by earlier
+            # prefetches we won't look for them here
+            discarded_from_prefetch_queue = set()
+            params_not_already_fetched = set(
+                filter(
+                    lambda p: self.__most_recent_step_id_param_fetched_for[p] < self.
+                    __step_id,
+                    params_to_fetch))
+            while self.__param_queue and len(discarded_from_prefetch_queue) < len(
+                    params_not_already_fetched):
+                param_in_trace = self.__param_queue.popleft()
+                self.__most_recent_step_id_param_fetched_for[
+                    param_in_trace.param] = param_in_trace.step_id_last_used_at
+                discarded_from_prefetch_queue.add(param_in_trace.param)
+            if discarded_from_prefetch_queue != params_not_already_fetched:
+                raise RuntimeError(
+                    f"tracing error at step {self.__step_id}: \n"
+                    f"module id: {current_submodule.id}, training: {current_submodule.training}\n"
+                    f"expected the next {len(params_not_already_fetched)} parameters in the "
+                    f"parameter fetch queue to be {tuple(p.ds_summary() for p in params_not_already_fetched)} \n"
+                    f"but got \n {tuple(p.ds_summary() for p in discarded_from_prefetch_queue)}."
+                )
+
+            # kick off all gather for params in the next few submodules (prefetch)
+            max_params_to_prefetch = min(
+                self.__max_n_available_params - self.__n_available_params,
+                self.__prefetch_bucket_sz)
+            params_to_prefetch = set()
+            numel_prefetching = 0
+            while self.__param_queue and numel_prefetching < max_params_to_prefetch:
+                param_in_trace: __class__.__ParamInTrace = self.__param_queue.popleft()
+                self.__most_recent_step_id_param_fetched_for[
+                    param_in_trace.param] = param_in_trace.step_id_last_used_at
+                if param_in_trace.param not in params_to_prefetch:
+                    params_to_prefetch.add(param_in_trace.param)
+                    numel_prefetching += param_in_trace.param.ds_numel
+            for param in params_to_prefetch:
+                debug_rank0(f"-prefetch: {param.ds_summary()}")
+            self.__all_gather_params(params_to_prefetch)
+
+            if self.__prefetch_nvme:
+                self.__prefetch_nvme_param_partitions()
+
+        self.__step_id += 1
+
+    @instrument_w_nvtx
+    @torch.no_grad()
+    def release_sub_module(self, submodule: Module) -> None:
+        """release the parameters of a sub module, assuming they meet conditions to
+        be released."""
+        params_to_release = (self.__params_to_release(submodule,
+                                                      self.__step_id)
+                             if self.trace_complete else set(
+                                 p.ds_id for p in iter_params(submodule)))
+
+        for param in iter_params(submodule):
+            param.ds_active_sub_modules.discard(submodule.id)
+            if param.ds_id in params_to_release and not param.is_external_param:
+                self.__release_param(param)
+
+    @instrument_w_nvtx
+    @torch.no_grad()
+    def release_and_reset_all(self, module: Module) -> None:
+        """release all module parameters"""
+        for param in iter_params(module, recurse=True):
+            if param in self.__inflight_param_registry:
+                raise RuntimeError(f"param {param.ds_summary()} still in flight")
+
+            # TODO. make this throw if if there are still active submodules. currently
+            # there's a hook execution issue
+            param.ds_active_sub_modules.clear()
+            self.__release_param(param)
+
+        for param in iter_params(module, recurse=True):
+            if param.ds_status != ZeroParamStatus.NOT_AVAILABLE:
+                raise RuntimeError(f"{param.ds_summary()} expected to be released")
+
+    @instrument_w_nvtx
+    def __all_gather_params(self, params: Set[Parameter]) -> None:
+        """for each partitioned parameter, kick off an async allgather and store
+        the work handle for the in flight parameters."""
+        partitioned_params = []
+        for param in params:
+            if param.ds_status == ZeroParamStatus.NOT_AVAILABLE:
+                partitioned_params.append(param)
+                self.__n_available_params += param.ds_numel
+
+        if partitioned_params:
+            with torch.cuda.stream(self.__allgather_stream):
+                handle = partitioned_params[0].all_gather_coalesced(partitioned_params)
+
+            for param in partitioned_params:
+                assert param.ds_status == ZeroParamStatus.INFLIGHT, param.ds_summary()
+                self.__inflight_param_registry[param] = handle
+
+    @instrument_w_nvtx
+    def __release_param(self, param: Parameter) -> None:
+        if param.ds_status == ZeroParamStatus.AVAILABLE and not param.ds_active_sub_modules:
+            debug_rank0(f"-release: {param.ds_summary()}")
+            param.partition()
+            self.__n_available_params -= param.ds_numel
+
+    @instrument_w_nvtx
+    @functools.lru_cache(maxsize=None)
+    def __params_to_release(self,
+                            submodule_to_release: Module,
+                            step_id: int) -> Set[int]:
+        if not self.trace_complete:
+            raise RuntimeError("expected trace to be complete")
+
+        params_to_release = set(p.ds_id for p in iter_params(submodule_to_release)
+                                if not p.ds_persist)
+
+        # examine all modules within `max_reuse_dist_in_numel` of the current step,
+        # if we see any of the candidate parameters to be released reoccur while
+        # doing this, remove them from the set of parameters to release.
+        params_traversed = 0
+        for module in self.__submodule_order[step_id:]:
+            if params_traversed > self.__max_reuse_dist_in_numel:
+                break
+            for param in iter_params(module):
+                params_to_release.discard(param.ds_id)
+                params_traversed += param.ds_numel
+
+        return params_to_release
+
+    @instrument_w_nvtx
+    def __prefetch_nvme_param_partitions(self) -> None:
+        """swap in parameter partitions from nvme for those parameters that will be used
+        after the ones that are already being prefetched into full parameters
+        """
+        if not self.trace_complete:
+            return
+
+        numel_in_flight = sum(param.ds_numel for param in self.__inflight_param_registry)
+
+        numel_considered = 0
+        swap_in_params = []
+        for param_in_trace in self.__param_queue:
+            param = param_in_trace.param
+            if param.nvme_swapper is None:
+                continue
+            if (numel_considered > 2 * numel_in_flight or len(swap_in_params) >=
+                    param.nvme_swapper.available_swap_in_buffers()):
+                break
+            if param.ds_tensor.status == PartitionedParamStatus.NOT_AVAILABLE:
+                swap_in_params.append(param)
+            numel_considered += param.ds_numel
+
+        if swap_in_params:
+            swap_in_params[0].nvme_swapper.swap_in(swap_in_params, async_op=True)
diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py
--- a/deepspeed/runtime/zero/stage3.py
+++ b/deepspeed/runtime/zero/stage3.py
@@ -30,6 +30,7 @@
 from deepspeed.ops.adam import DeepSpeedCPUAdam
 from deepspeed.ops.op_builder import UtilsBuilder
 from deepspeed.runtime.zero.offload_constants import *
+from deepspeed.runtime.zero.partitioned_param_coordinator import PartitionedParameterCoordinator, iter_params
 from deepspeed.runtime.swap_tensor.partitioned_param_swapper import PartitionedParamStatus
 from deepspeed.runtime.swap_tensor.partitioned_optimizer_swapper import PartitionedOptimizerSwapper
 from deepspeed.runtime.swap_tensor.pipelined_optimizer_swapper import PipelinedOptimizerSwapper
@@ -67,26 +68,11 @@ def lcm(x, y):
     return x * y // gcd(x, y)
 
 
-def debug_rank0(message: str) -> None:
-    if dist.get_rank() == 0:
-        logger.debug(message)
-
-
 def move_to_cpu(tensor_list):
     for tensor in tensor_list:
         tensor.data = tensor.data.cpu()
 
 
-@instrument_w_nvtx
-def get_all_parameters(sub_module, recurse=False):
-    return itertools.chain(sub_module.named_parameters(recurse=recurse),
-                           sub_module.ds_external_parameters())
-
-
-def iter_params(module: Module, recurse=False) -> Iterable[Parameter]:
-    return map(lambda pair: pair[1], get_all_parameters(module, recurse))
-
-
 #apply torch.autograd.Function that calls a backward_function to tensors in output
 def _apply_to_tensors_only(module, functional, backward_function, outputs):
     if type(outputs) is tuple:
@@ -149,10 +135,11 @@ def __getitem__(self, key):
 
         if param.ds_status == ZeroParamStatus.NOT_AVAILABLE:
             if self._parent_module._parameters._in_forward:
-                print_rank_0(f'Registering external parameter from getter {key}',
-                             force=False)
                 register_external_parameter(FWD_MODULE_STACK[-1], param)
                 param.all_gather()
+                print_rank_0(
+                    f'Registering external parameter from getter {key} ds_id = {param.ds_id}',
+                    force=False)
 
         return param
 
@@ -169,329 +156,6 @@ def _inject_parameters(module, cls):
         module._parameters = new_param
 
 
-class PartitionedParameterCoordinator:
-    """Handles partitioning and gathering of parameters."""
-    class __InflightParamRegistry(UserDict):
-        """registry for parameters in flight"""
-        def __setitem__(self,
-                        param: Parameter,
-                        handle: AllGatherCoalescedHandle) -> None:
-            if param in self.data:
-                raise RuntimeError(f"{param.ds_summary()} already in registry")
-            if param.ds_status != ZeroParamStatus.INFLIGHT:
-                raise RuntimeError(
-                    f"attempted to add non-inflight parameter to registry {param.ds_summary()}"
-                )
-            self.data[param] = handle
-
-    @dataclass
-    class __ParamInTrace:
-        param: Parameter
-        step_id_last_used_at: int
-
-    def __init__(
-        self,
-        prefetch_bucket_sz: int,
-        max_reuse_distance_in_numel: int,
-        max_available_parameters_in_numel: int,
-        allgather_stream: Stream,
-        prefetch_nvme: bool = False,
-    ) -> None:
-        # mapping of param -> handle for each param that is currently in flight
-        self.__inflight_param_registry = __class__.__InflightParamRegistry()
-        # keeps track of the number of submodules invoked so far.
-        self.__step_id: int = 0
-        # whether or not we have completed a trace of the entire network. This should
-        # always be true after the first forward pass + backward pass.
-        self.trace_complete: bool = False
-        # sequence of submodules/parameters in forward pass + backward pass
-        self.__submodule_order: Iterable[Module] = []
-        self.__param_order: Iterable[__class__.__ParamInTrace] = []
-        self.__most_recent_step_id_param_fetched_for = collections.defaultdict(
-            lambda: int(-1e10))
-        # number of available params, and max number of available params
-        self.__n_available_params: int = 0
-        self.__max_n_available_params: int = max_available_parameters_in_numel
-        # max distance between two use of the module beyond which module is released
-        self.__max_reuse_dist_in_numel: int = max_reuse_distance_in_numel
-        # queue for parameters to fetch. parameters will be popped off the left
-        # side of the dequeue as they are fetched
-        self.__param_queue: Deque[__class__.__ParamInTrace] = None
-        self.__prefetch_bucket_sz: int = prefetch_bucket_sz
-        self.__prefetch_nvme: bool = prefetch_nvme
-        self.hierarchy: int = 0
-
-        # stream that will be used for allgather operations
-        self.__allgather_stream: Stream = allgather_stream
-
-        # limit the number of fetch events that can be queued at once
-        # otherwise, what happens is memory is allocated by the host thread at the
-        # time of the call, but not used until later by the asynchronous cuda stream.
-        # allowing an infinite number of these to queue up causes a lot of memory
-        # pressure that then becomes detrimental to performance.
-        # this is a much less elegant way of fixing this vs something like using
-        # cudaMallocAsync/cudaFreeAsync. Choosing to not expose this to the user now
-        # because ideally in the future its replaced by an async allocation
-        # mechanism which doesn't require any configuration by the user.
-        self.__ongoing_fetch_events: Deque[Event] = collections.deque()
-        # TODO. make this configurable via JSON
-        self.__max_ongoing_fetch_events: int = 2
-
-    """Tracing and Tracking
-    TODO. consider performing trace before initializing PartitionedParameterCoordinator
-    and passing trace results into constructor. This way all the code in here can
-    just assume that the trace is complete and the results can be entirely
-    immutable.
-
-    Bookkeeping operations used to track where we are in the forward/backward pass
-    """
-
-    def record_trace(self, sub_module: Module) -> None:
-        """adds sub module to trace"""
-        if self.trace_complete:
-            raise RuntimeError(
-                "attempted to record trace when trace was already complete")
-
-        self.__submodule_order.append(sub_module)
-        for param in sorted(set(iter_params(sub_module)), key=lambda p: p.ds_id):
-            self.__param_order.append(
-                __class__.__ParamInTrace(param=param,
-                                         step_id_last_used_at=self.__step_id))
-
-    def reset_step(self) -> None:
-        """indicate that we have completed one fwd+bwd for the model"""
-        if self.__inflight_param_registry:
-            raise RuntimeError(
-                f"still have inflight params "
-                f"{[p.ds_summary for p in self.__inflight_param_registry.keys()]}")
-
-        if not self.trace_complete:
-            # make sure that recorded parameter and submodule orders are
-            # identical across ranks
-            assert_ints_same_as_other_ranks([m.id for m in self.__submodule_order])
-            assert_ints_same_as_other_ranks([p.param.ds_id for p in self.__param_order])
-            assert_ints_same_as_other_ranks(
-                [p.step_id_last_used_at for p in self.__param_order])
-
-            self.__submodule_order = tuple(self.__submodule_order)  # freeze
-            self.__param_order = tuple(self.__param_order)  # freeze
-            self.trace_complete = True
-            print_rank_0(f"completed trace: {[m.id for m in self.__submodule_order]}",
-                         force=False)
-
-        self.__param_queue = collections.deque(self.__param_order)  # reset fetch queue
-        self.__most_recent_step_id_param_fetched_for = collections.defaultdict(
-            lambda: int(-1e10))
-        self.__step_id = 0
-        self.__n_available_params = 0
-
-    """Fetch and Release
-    Fetching, prefetching, and releasing parameters
-    """
-
-    @instrument_w_nvtx
-    @torch.no_grad()
-    def fetch_sub_module(self, current_submodule: Module) -> None:
-        """This method does the following (in order):
-        1. kick off fetch for parameters in immediately required sub module
-        2. kick off fetch for next few parameters we will need later (prefetch)
-        3. block on parameters in immediately required sub module
-        """
-        debug_rank0(
-            f"{self.__step_id}: M{current_submodule.id}({type(current_submodule).__name__}) P{[p.ds_id for p in iter_params(current_submodule)]} "
-            + str({
-                "avail": f"{self.__n_available_params:.1e}",
-                "queue_sz": f"{len(self.__param_queue or [])}",
-                "inflight": [p.ds_id for p in self.__inflight_param_registry],
-            }))
-
-        params_to_fetch = frozenset(iter_params(current_submodule))
-
-        # kick off all gather for params in the immediately required submodule
-        for param in params_to_fetch:
-            debug_rank0(f"-fetch: {param.ds_summary()}")
-        self.__all_gather_params(params_to_fetch)
-
-        # wait for parameters in the immediately needed submodule to become available
-        for param in params_to_fetch:
-            param.ds_active_sub_modules.add(current_submodule.id)
-            debug_rank0(f"-wait: {param.ds_summary()}")
-            if param in self.__inflight_param_registry:
-                with torch.cuda.stream(self.__allgather_stream):
-                    while self.__ongoing_fetch_events and self.__ongoing_fetch_events[
-                            0].query():
-                        self.__ongoing_fetch_events.popleft()
-                    if len(self.__ongoing_fetch_events
-                           ) > self.__max_ongoing_fetch_events:
-                        self.__ongoing_fetch_events.popleft().synchronize()
-
-                    self.__inflight_param_registry.pop(param).wait()
-
-                    event = Event()
-                    event.record()
-                    self.__ongoing_fetch_events.append(event)
-
-            assert param.ds_status == ZeroParamStatus.AVAILABLE, param.ds_summary()
-        torch.cuda.current_stream().wait_stream(self.__allgather_stream)
-
-        # kick off parameter prefetches for upcoming modules
-        # don't prefetch if we dont have a completed model trace, or if we aren't
-        # training (throws off the tracing and don't want to prefetch modules for bwd)
-        if self.trace_complete and current_submodule.training:
-            # go through the parameters we need for the current module and pop them
-            # off the fetch queue so that they aren't prefetched later.
-            # if params have already been popped off the fetch queue by earlier
-            # prefetches we won't look for them here
-            discarded_from_prefetch_queue = set()
-            params_not_already_fetched = set(
-                filter(
-                    lambda p: self.__most_recent_step_id_param_fetched_for[p] < self.
-                    __step_id,
-                    params_to_fetch))
-            while self.__param_queue and len(discarded_from_prefetch_queue) < len(
-                    params_not_already_fetched):
-                param_in_trace = self.__param_queue.popleft()
-                self.__most_recent_step_id_param_fetched_for[
-                    param_in_trace.param] = param_in_trace.step_id_last_used_at
-                discarded_from_prefetch_queue.add(param_in_trace.param)
-            if discarded_from_prefetch_queue != params_not_already_fetched:
-                raise RuntimeError(
-                    f"tracing error at step {self.__step_id}: "
-                    f"expected the next {len(params_not_already_fetched)} parameters in the "
-                    f"parameter fetch queue to be {tuple(p.ds_summary() for p in params_not_already_fetched)} "
-                    f"but got {tuple(p.ds_summary() for p in discarded_from_prefetch_queue)}."
-                )
-
-            # kick off all gather for params in the next few submodules (prefetch)
-            max_params_to_prefetch = min(
-                self.__max_n_available_params - self.__n_available_params,
-                self.__prefetch_bucket_sz)
-            params_to_prefetch = set()
-            numel_prefetching = 0
-            while self.__param_queue and numel_prefetching < max_params_to_prefetch:
-                param_in_trace: __class__.__ParamInTrace = self.__param_queue.popleft()
-                self.__most_recent_step_id_param_fetched_for[
-                    param_in_trace.param] = param_in_trace.step_id_last_used_at
-                if param_in_trace.param not in params_to_prefetch:
-                    params_to_prefetch.add(param_in_trace.param)
-                    numel_prefetching += param_in_trace.param.ds_numel
-            for param in params_to_prefetch:
-                debug_rank0(f"-prefetch: {param.ds_summary()}")
-            self.__all_gather_params(params_to_prefetch)
-
-            if self.__prefetch_nvme:
-                self.__prefetch_nvme_param_partitions()
-
-        self.__step_id += 1
-
-    @instrument_w_nvtx
-    @torch.no_grad()
-    def release_sub_module(self, submodule: Module) -> None:
-        """release the parameters of a sub module, assuming they meet conditions to
-        be released."""
-        params_to_release = (self.__params_to_release(submodule,
-                                                      self.__step_id)
-                             if self.trace_complete else set(
-                                 p.ds_id for p in iter_params(submodule)))
-
-        for param in iter_params(submodule):
-            param.ds_active_sub_modules.discard(submodule.id)
-            if param.ds_id in params_to_release and not param.is_external_param:
-                self.__release_param(param)
-
-    @instrument_w_nvtx
-    @torch.no_grad()
-    def release_and_reset_all(self, module: Module) -> None:
-        """release all module parameters"""
-        for param in iter_params(module, recurse=True):
-            if param in self.__inflight_param_registry:
-                raise RuntimeError(f"param {param.ds_summary()} still in flight")
-
-            # TODO. make this throw if if there are still active submodules. currently
-            # there's a hook execution issue
-            param.ds_active_sub_modules.clear()
-            self.__release_param(param)
-
-        for param in iter_params(module, recurse=True):
-            if param.ds_status != ZeroParamStatus.NOT_AVAILABLE:
-                raise RuntimeError(f"{param.ds_summary()} expected to be released")
-
-    @instrument_w_nvtx
-    def __all_gather_params(self, params: Set[Parameter]) -> None:
-        """for each partitioned parameter, kick off an async allgather and store
-        the work handle for the in flight parameters."""
-        partitioned_params = []
-        for param in params:
-            if param.ds_status == ZeroParamStatus.NOT_AVAILABLE:
-                partitioned_params.append(param)
-                self.__n_available_params += param.ds_numel
-
-        if partitioned_params:
-            with torch.cuda.stream(self.__allgather_stream):
-                handle = partitioned_params[0].all_gather_coalesced(partitioned_params)
-
-            for param in partitioned_params:
-                assert param.ds_status == ZeroParamStatus.INFLIGHT, param.ds_summary()
-                self.__inflight_param_registry[param] = handle
-
-    @instrument_w_nvtx
-    def __release_param(self, param: Parameter) -> None:
-        if param.ds_status == ZeroParamStatus.AVAILABLE and not param.ds_active_sub_modules:
-            debug_rank0(f"-release: {param.ds_summary()}")
-            param.partition()
-            self.__n_available_params -= param.ds_numel
-
-    @instrument_w_nvtx
-    @functools.lru_cache(maxsize=None)
-    def __params_to_release(self,
-                            submodule_to_release: Module,
-                            step_id: int) -> Set[int]:
-        if not self.trace_complete:
-            raise RuntimeError("expected trace to be complete")
-
-        params_to_release = set(p.ds_id for p in iter_params(submodule_to_release)
-                                if not p.ds_persist)
-
-        # examine all modules within `max_reuse_dist_in_numel` of the current step,
-        # if we see any of the candidate parameters to be released reoccur while
-        # doing this, remove them from the set of parameters to release.
-        params_traversed = 0
-        for module in self.__submodule_order[step_id:]:
-            if params_traversed > self.__max_reuse_dist_in_numel:
-                break
-            for param in iter_params(module):
-                params_to_release.discard(param.ds_id)
-                params_traversed += param.ds_numel
-
-        return params_to_release
-
-    @instrument_w_nvtx
-    def __prefetch_nvme_param_partitions(self) -> None:
-        """swap in parameter partitions from nvme for those parameters that will be used
-        after the ones that are already being prefetched into full parameters
-        """
-        if not self.trace_complete:
-            return
-
-        numel_in_flight = sum(param.ds_numel for param in self.__inflight_param_registry)
-
-        numel_considered = 0
-        swap_in_params = []
-        for param_in_trace in self.__param_queue:
-            param = param_in_trace.param
-            if param.nvme_swapper is None:
-                continue
-            if (numel_considered > 2 * numel_in_flight or len(swap_in_params) >=
-                    param.nvme_swapper.available_swap_in_buffers()):
-                break
-            if param.ds_tensor.status == PartitionedParamStatus.NOT_AVAILABLE:
-                swap_in_params.append(param)
-            numel_considered += param.ds_numel
-
-        if swap_in_params:
-            swap_in_params[0].nvme_swapper.swap_in(swap_in_params, async_op=True)
-
-
 class PreBackwardFunction(torch.autograd.Function):
     @staticmethod
     def forward(ctx, module, pre_backward_function, outputs):
@@ -655,18 +319,14 @@ def __init__(self,
         ############################################################################
 
         see_memory_usage("Before Partitioned Parameter Coordinator", force=False)
-        self.param_coordinator = PartitionedParameterCoordinator(
-            prefetch_bucket_sz=int(prefetch_bucket_size),
-            max_reuse_distance_in_numel=int(max_reuse_distance),
-            max_available_parameters_in_numel=int(max_live_parameters),
-            allgather_stream=self.__allgather_stream,
-            prefetch_nvme=self.params_in_nvme_and_cpu,
-        )
+        self.param_coordinators = {}
+        self._prefetch_bucket_sz = int(prefetch_bucket_size)
+        self._max_reuse_distance_in_numel = int(max_reuse_distance)
+        self._max_available_parameters_in_numel = int(max_live_parameters)
         see_memory_usage("After Partitioned Parameter Coordinator", force=False)
 
         self.__n_caching_allocator_flushes = 0
 
-        #self.param_coordinator = PartitionedParameterCoordinator(comm_stream=torch.cuda.Stream())
         #-------------Stage 3 Setup-------------------#
         # parameters smaller than the threshold will be collectively gathered at the
         # end of the optimizer step and will be kept till the end of the backward pass
@@ -934,6 +594,19 @@ def defragment(tensors: List[Tensor]) -> Tensor:
 
         return device_buffer
 
+    def _get_param_coordinator(self, training):
+        if not training in self.param_coordinators:
+            self.param_coordinators[training] = PartitionedParameterCoordinator(
+                prefetch_bucket_sz=self._prefetch_bucket_sz,
+                max_reuse_distance_in_numel=self._max_reuse_distance_in_numel,
+                max_available_parameters_in_numel=self.
+                _max_available_parameters_in_numel,
+                allgather_stream=self.__allgather_stream,
+                prefetch_nvme=self.params_in_nvme_and_cpu,
+            )
+
+        return self.param_coordinators[training]
+
     def _configure_offloading(self, offload_optimizer_config, offload_param_config):
         ###################### offload optimizer setup ##################################
         if offload_optimizer_config is not None:
@@ -1370,7 +1043,7 @@ def setup_zero_stage3_hooks(self):
         def _end_of_forward_hook(module, *args):
 
             if not torch._C.is_grad_enabled():
-                self.param_coordinator.reset_step()
+                self._get_param_coordinator(training=False).reset_step()
 
         #likely one of them should be enough but just to be safe
         self._register_hooks_recursively(self.module)
@@ -1433,16 +1106,16 @@ def _post_forward_module_hook(module, input, output):
                 if not any(id(item) in m._external_params for m in FWD_MODULE_STACK):
                     item.is_external_param = True
                     module_to_register = FWD_MODULE_STACK[-1]
+                    register_external_parameter(module_to_register, item)
                     print_rank_0(
-                        f'Registering dangling parameter for module {module_to_register.__class__.__name__}.',
+                        f'Registering dangling parameter for module {module_to_register.__class__.__name__}, ds_id = {item.ds_id}.',
                         force=False)
-                    register_external_parameter(module_to_register, item)
 
                     # It's possible that the parameter was already external to the completed module. If so, remove it the
                     # registration as it will be covered by the outer module instead.
                     if id(item) in module._external_params:
                         print_rank_0(
-                            f'  Unregistering nested dangling parameter from module {module.__class__.__name__}',
+                            f'  Unregistering nested dangling parameter from module {module.__class__.__name__}, ds_id = {item.ds_id}',
                             force=False)
                         unregister_external_parameter(module, item)
 
@@ -1522,10 +1195,8 @@ def pre_sub_module_forward_function(self, sub_module):
         global FWD_MODULE_STACK
         FWD_MODULE_STACK.append(sub_module)
 
-        if not self.param_coordinator.trace_complete:
-            self.param_coordinator.record_trace(sub_module)
-
-        self.param_coordinator.fetch_sub_module(sub_module)
+        self._get_param_coordinator(
+            training=sub_module.training).fetch_sub_module(sub_module)
         see_memory_usage(
             f"Before sub module function {sub_module.__class__.__name__} after fetch",
             force=False)
@@ -1536,7 +1207,11 @@ def post_sub_module_forward_function(self, sub_module):
             f"After sub module function {sub_module.__class__.__name__} {sub_module.id} before release",
             force=False)
 
-        self.param_coordinator.release_sub_module(sub_module)
+        param_coordinator = self._get_param_coordinator(training=sub_module.training)
+        if not param_coordinator.trace_complete:
+            param_coordinator.record_trace(sub_module)
+
+        param_coordinator.release_sub_module(sub_module)
 
         see_memory_usage(
             f"After sub module function {sub_module.__class__.__name__}  {sub_module.id} after release",
@@ -1544,16 +1219,20 @@ def post_sub_module_forward_function(self, sub_module):
 
     @torch.no_grad()
     def pre_sub_module_backward_function(self, sub_module):
-        if not self.param_coordinator.trace_complete:
-            self.param_coordinator.record_trace(sub_module)
-        self.param_coordinator.fetch_sub_module(sub_module)
+        param_coordinator = self._get_param_coordinator(training=sub_module.training)
+        if not param_coordinator.trace_complete:
+            param_coordinator.record_trace(sub_module)
+        param_coordinator.fetch_sub_module(sub_module)
 
     @torch.no_grad()
     def post_sub_module_backward_function(self, sub_module):
         see_memory_usage(
             f"After sub module backward function {sub_module.__class__.__name__} {sub_module.id} before release",
             force=False)
-        self.param_coordinator.release_sub_module(sub_module)
+
+        self._get_param_coordinator(
+            training=sub_module.training).release_sub_module(sub_module)
+
         see_memory_usage(
             f"After sub module backward function {sub_module.__class__.__name__} {sub_module.id} after release",
             force=False)
@@ -2400,7 +2079,7 @@ def _pre_step(self):
         see_memory_usage(f"In step before checking overflow", force=False)
 
         print_rank_0("Finished Tracing at Beginning of Step")
-        self.param_coordinator.hierarchy = 0
+        self._get_param_coordinator(training=True).hierarchy = 0
 
         print_rank_0("Finished Tracing at Beginning of Step")
 
@@ -2795,7 +2474,7 @@ def backward(self, loss, retain_graph=False):
 
         self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
 
-        self.param_coordinator.reset_step()
+        self._get_param_coordinator(training=True).reset_step()
 
         if self.swap_optimizer:
             self.optimizer_swapper.post_backward()
@@ -2827,7 +2506,8 @@ def _partition_all_parameters(self):
         """Partitioning Parameters that were not partitioned usually if parameters
         of modules whose input parameters do not require grad computation do not
         trigger post call and will therefore will remain unpartitioned"""
-        self.param_coordinator.release_and_reset_all(self.module)
+        self._get_param_coordinator(training=self.module.training).release_and_reset_all(
+            self.module)
         for param in iter_params(self.module, recurse=True):
             if param.ds_status != ZeroParamStatus.NOT_AVAILABLE:
                 raise RuntimeError(f"{param.ds_summary()} expected to be released")


Test patch: