[Training] Add an extensive set of optimizer tests

[vladimirjovanovicTT]

We want to make sure optimizer functionality works for the following: optimizer in vacuum, using a toy model+optimizer pass, and using sth resembling a real-world training scenario.

Proposed testing setup is the following:

1. Test internal optimizer functionality:

• Using hardcoded inputs and states, comparing new optimizer outputs/states to golden.

2. Test optimizer functionality in a training scenario:

• Hardcode weights and grads; do optimizer step, run a fwd pass with updated weights and compare model outputs to golden.

3. test e2e

• Initialize weights in a standard way (xavier or similar), run fwd/bwd passes through model, do optimizer step and check results for the following fwd pass.

4. We might want to make sure that 3 is stable through a sequence of minibatches (seq length is a hyperparam).

We can create sub-issues for each of these, if necessary.