Automatic Stochastic depth on residual blocks

[dskhudia]

Detect residual blocks automatically and replace them with a stochastic version. See attached graphs for how transformation happens on ResNet50.

TODO: Add tests. Done

Orig:

Split into bottleneck and downsample:

Converted with a call to block_stochastic_module:

Closes https://mosaicml.atlassian.net/browse/CO-303

[Landanjs]

This is awesome! It LGTM from the stochastic depth side. A few nits up to your discretion

I'm a lot less confident on the fx graph surgery implementation. It looks reasonable to me, but not sure if there are alternative implementations? It would be great to get someone from eng to review this as well.

One hesitation before approving: would it be possible to create a test for apply_stochastic_residual? It be nice to check that the graph surgery went as expected for say ResNet18 and ResNet50 e.g. counting the number of stochastic blocks after surgery and checking it aligns with the number of residual blocks before surgery

[Landanjs]

Oh, you just pushed a test 😂 Is this still being updated?

[dskhudia]

Oh, you just pushed a test 😂 Is this still being updated?

Just the test was missing so added it a little while ago.

[dskhudia]

It be nice to check that the graph surgery went as expected for say ResNet18 and ResNet50 e.g. counting the number of stochastic blocks after surgery and checking it aligns with the number of residual blocks before surgery

Added a test for resnet18 that counts residual blocks as well as runs it with 0.0 drop_rate to make sure pre and post surgery outputs are the same.

[ravi-mosaicml]

Going to defer the correctness review to @Landanjs / @bandish-shah, but see comments re: code styling.

[hanlint]

Nice work, a few style nits, and a main few questions in the comments.

[ravi-mosaicml]

Overall looks good, thanks for doing this. See comments. I didn't check for correctness FYI; maybe wait for @hanlint or @Landanjs to approve before merging?

[ravi-mosaicml]

(Dismissing my review; don't have much to add)

[hanlint]

LGTM, one small nit. Defer to @Landanjs to review for correctness.

Right now apply_stochastic_residual is not actually used in the stochastic depth algorithm. I assume this will be deferred to follow-on PR? Should we do a training run with this to see if matches the performance of the old apply?

[dskhudia]

LGTM, one small nit. Defer to @Landanjs to review for correctness.

Right now apply_stochastic_residual is not actually used in the stochastic depth algorithm. I assume this will be deferred to follow-on PR? Should we do a training run with this to see if matches the performance of the old apply?

It will be applied in a followup PR along with experiments. However, I did a training run with r50 baseline (77.32 Top-1 in 132.8 mins) and r50 baseline with fx-traced module (77.25 in 125.4 mins). We shouldn't read too much into the times here as it could be due to streaming datasets. I don't see much issues when this is applied with stochastic depth algorithm but, as always, composition with others might pose some issues.

[Landanjs]

The stochastic depth implementation LGTM. Like Hanlin mentioned the best test would be before and after experiments, but those will happen in a follow up PR 🙂