Figure out why profiler breaks after running

[StellaAthena]

We can profile the code with with examples/ds_pretrain_gpt2_pipe_flops_profiler.sh. It seems to break after profiling, but still returns the profile result. Not sure what's going on here and it's worth looking into. Example traceback:

  File "pretrain_gpt2.py", line 156, in <module>
    pretrain(train_valid_test_datasets_provider, model_provider, forward_step,
  File "/root/megatron-3d/megatron/training.py", line 97, in pretrain
    iteration = train(forward_step_func,
  File "/root/megatron-3d/megatron/training.py", line 481, in train
    loss_dict, skipped_iter = train_step(forward_step_func,
  File "/root/megatron-3d/megatron/training.py", line 324, in train_step
    return train_step_pipe(model, data_iterator)
  File "/root/megatron-3d/megatron/training.py", line 358, in train_step_pipe
    loss = model.train_batch(data_iter=data_iterator)
  File "/root/anaconda3/lib/python3.8/site-packages/deepspeed/runtime/pipe/engine.py", line 283, in train_batch
    self._exec_schedule(sched)
  File "/root/anaconda3/lib/python3.8/site-packages/deepspeed/runtime/pipe/engine.py", line 1172, in _exec_schedule
    self._exec_instr(**cmd.kwargs)
  File "/root/anaconda3/lib/python3.8/site-packages/deepspeed/runtime/pipe/engine.py", line 521, in _exec_forward_pass
    outputs = super().forward(inputs)
  File "/root/anaconda3/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 824, in forward
    loss = self.module(*inputs, **kwargs)
  File "/root/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 719, in _call_impl
    result = hook(self, input)
  File "/root/anaconda3/lib/python3.8/site-packages/deepspeed/profiling/flops_profiler/profiler.py", line 49, in pre_hook
    module.__steps__ += 1
  File "/root/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 778, in __getattr__
    raise ModuleAttributeError("'{}' object has no attribute '{}'".format(
torch.nn.modules.module.ModuleAttributeError: 'GPT2ModelPipe' object has no attribute '__steps__'```

[sdtblck]

probably worth an issue on Deepspeed