Improve and refine MLP tests for extensibility and A/B testing

[rpsilva-aws]

In this PR, we include various fixes, improvements and extensions, namely:

• Exposing the MLP test to other tests (to allow us to A/B test convergences: requirement for the grad acc tests)
• Improving the asserts, and extending to losses and outputs
• Using the appropriate flag for the LR and log steps
• Improving the model layout with nn.Sequential
• Enhance the coverage, by actually utilizing the checkpointing flag, and including a sanity test for CPU
• Decouple to simplify the A/B coverage
• Fix imports

[rpsilva-aws]

@tengyifei @ManfeiBai I found myself having to largely improve/enhance the MLP tests, since I wanted to reuse this test for A/B convergence validation:

• Gradient accumulation without the loop (batch size = 128)
• Gradient accumulation with the loop (experimental grad acc) (batch size = 128)
• Higher DP degree, no grad acc (batch size = 128)

PTAL.

[tengyifei]

Synced offline. Let's add a --skip-gradient-checkpointing CLI arg to the train testing script or similar to skip the gradient checkpointing on CPU, in order to avoid the confusing test_train_spmd_linear_model_grad_checkpointing name.

[rpsilva-aws]

Synced offline. Let's add a --skip-gradient-checkpointing CLI arg to the train testing script or similar to skip the gradient checkpointing on CPU, in order to avoid the confusing test_train_spmd_linear_model_grad_checkpointing name.

I was just about to write this, done!

[lsy323]

Looks like the test file is duplicated under test/utils?

[rpsilva-aws]

Looks like the test file is duplicated under test/utils?

I don't see it. The core of the functionality was moved to utils - it's a standalone training script. The test_ file can run simultaneous executions of it, to, intentionally, run A/B testing. The test/utils file is not for direct testing, but users can run it as a standalone script, and so can tests.

[lsy323]

Looks like the test file is duplicated under test/utils?

I don't see it. The core of the functionality was moved to utils - it's a standalone training script. The test_ file can run simultaneous executions of it, to, intentionally, run A/B testing. The test/utils file is not for direct testing, but users can run it as a standalone script, and so can tests.

Hi @rpsilva-aws thank you for the explanation. I still don't get the A/B testing part. But if the body can be reused in other tests then it's fine

[rpsilva-aws]

Looks like the test file is duplicated under test/utils?

I don't see it. The core of the functionality was moved to utils - it's a standalone training script. The test_ file can run simultaneous executions of it, to, intentionally, run A/B testing. The test/utils file is not for direct testing, but users can run it as a standalone script, and so can tests.

Hi @rpsilva-aws thank you for the explanation. I still don't get the A/B testing part. But if the body can be reused in other tests then it's fine

Thanks. You can see it in this PR here: https://github.com/pytorch/xla/pull/8561/files#diff-09c6a280d1c6fc8053d5a17e919e29fe51a75fd593f6c2012496b6e970312c25R44

Essentially, it allows it to A/B test different functionality. It adds a bit more value than running each standalone test, and just checking that the losses/output are not 0. Alternatively, we'd need to assert against an hardcoded set of expected values. We want to reuse this for testing gradient accumulation with and without XLA's while loop.