Big science fix passing multiple tensors

[thomasw21]

Fixes necessary for implementation of prefix lm:

• Passing tuple of tensors between stages
• Allowing models to pass bool tensors. (This seems fixed in torch.distributed for pytorch > 1.7)

[stas00]

Great work, Thomas!

I left a few small suggestions.

I think the only issue here is that this new code doesn't have a test, so unless we add one it's likely to break down the road.

I'm not an expert on DS test suite, I added a few tests following monkey-see-monkey-do but this one might be not so easy to add. Let's first see if this is in general acceptable by the DS team and then we can work on a test.

[stas00]

should this not be a tuple? it could be just fine, I was just comparing with the original and it was there...

[thomasw21]

This is a tuple, in python you don't need to specifcy tuple, parenthesis work just fine.

[hyunwoongko]

I think it will be list.

>>> list_ = [1, 2, 3]
>>> is_tuple = (list_)
>>> type(is_tuple)
<class 'list'>

the following line means type casting. parenthesis can not cast list into tuple.
so you need to remain keyword tuple.

inputs = ([part.to_meta(), part.data()])  # list
inputs = tuple([part.to_meta(), part.data()])  # tuple

[thomasw21]

ah well played I didn't notice, which now begs the question why it works ...

EDIT: It actually doesn't matter, because you end up looping over the elements of an iterable. But I'll make the desired changes.

[stas00]

@ShadenSmith, @jeffra - could you please have a look at the proposed changes - this is currently a blocker for our adding of prefix-lm to Meg-DS. Thank you!

[stas00]

@thomasw21, also please run:

pre-commit run --all-files

to auto-format your code - as you can see the build CI failed.

pip install pre-commit 

[hyunwoongko]

[hyunwoongko]

@thomasw21 I'll test your branch now. If it works, this will be a very great work.

[hyunwoongko]

@thomasw21 I got error from this line.

assert all([torch.is_tensor(elt) and elt.requires_grad is False for elt in outputs[1:]])

Why all the output tensors except for first one must not be grad tensor when tuple was output?

[hyunwoongko]

I think this assertion torch.is_tensor() came from big-science branch. why did you guys change like it?

407ff0f#diff-26611f6be759237464a03bb1328cbc16555888836b3504dc3703e2e25d2a3ca3

torch.is_tensor() came from this PR.

[thomasw21]

Hey @hyunwoongko! Concerning the grad issue, i had two reasons:

• it doesn't fit our current need. Typically we wanted something similar to GPT2ModelPipe hacks. But in my opinion this is a more generic solution.
• I'm not super sure on how we'd want to handle that, typically with PartitionedTensor. There are a few tricks with tensors that require grads, which I'm a bit unfamiliar with. Typically the way grads are stored https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/pipe/engine.py#L595-L597 So ideally I guess we'd want a similar way of handling grads to reduce memory footprint no?

[hyunwoongko]

#1432

I made un issue to make it more accessible for others.
I think someone may have similar problems later on as we do.

[thomasw21]

@stas00 Thanks ! I'll wait to get first general approval of the approach before writing tests I guess?

Something important to note is that passing bool tensors works because we use torch 1.7+ (there was an issue with bool tensors not working in distributed setting using nccl)

• Discussion: Fix deepspeed prefix-lm bigscience-workshop/Megatron-DeepSpeed#107 (comment)
• Related issue (fixed): [distributed] NCCL Backend doesn't support torch.bool data type pytorch/pytorch#24137

[stas00]

@stas00 Thanks ! I'll wait to get first general approval of the approach before writing tests I guess?

Something important to note is that passing bool tensors works because we use torch 1.7+ (there was an issue with bool tensors not working in distributed setting using nccl)

perhaps then it should check that pytorch>=1.7 - I don't know if we could simply require that for the pipeline files?

@jeffra, what do you think? Do you have internal needs to support torch<1.7 for PP?

Additionally, could one of you please confirm that this line of changes is agreeable to you?

[thomasw21]

I should have fixed all the failing tests:

pytest --forked unit/test_checkpointing.py::test_checkpoint_moe unit/test_checkpointing.py::test_checkpoint_zero_no_optimizer unit/test_checkpointing.py::test_checkpoint_unique_tag unit/test_configurable_parallel.py::TestConfigurablePP::test_gpt2_mp2_pp_2to1 unit/test_configurable_parallel.py::TestConfigurablePP::test_gpt2_mp1_pp_2to1 unit/test_configurable_parallel.py::TestConfigurablePP::test_pp_basic unit/test_configurable_parallel.py::TestConfigurablePP::test_gpt2_pp_1to2_mp_1to2 unit/test_configurable_parallel.py::TestConfigurablePP::test_gpt2_mp2_pp_1to2 unit/test_configurable_parallel.py::TestConfigurablePP::test_gpt2_pp_2to1_mp_2to1 unit/test_configurable_parallel.py::TestConfigurableMP::test_gpt2_mp2_no_resize

[stas00]

@jeffra, @tjruwase - could you please start CI? Thank you!

[ShadenSmith]

This is great, thanks a ton!