[IPU] Support `self.log` within the step function

[SeanNaren]

🚀 Feature

As discussed in #7867 currently we do not support self.log in the step function as the function is traced. However with the help from the Graphcore engineer's on their side, it should be possible to expose tensors that do not have grads associated with them in the step function or after the step function.

Once this has been set in place, we should be able to support self.log for most common use cases.

I know work has already began on the Graphcore side cc @Borda @tchaton @SeanNaren @avors72 @AnthonyBarbier

[tchaton]

Dear @SeanNaren @avors72 @AnthonyBarbier,

Any estimation around the ETA of the logging feature ?

Best,
T.C

[AnthonyBarbier]

For regular tensors (i.e tensors from the forward pass), it should be possible to implement this using Anchors with PopTorch v2.1
Gradients / optimizer states will require much more work in the lower parts of our stack so I can't really give an estimate yet.

Please note this will have a significant impact in terms of both speed and memory consumption: so we'd recommend using it for debugging purposes but not for any run where performance matters.

[SeanNaren]

Thanks @AnthonyBarbier!

I think for the most part tensors that do not have gradients is the minimal requirement and fixes a host of logging issues.

You mention speed impacts, do you know a rough estimate on how much of an impact these anchors will make?

We might be able to get away with adding logic to warn the user of the speed hit + allow disabling of logging to make it a no-op.

[davors72]

The problem is a bit hard to quantify in cycles in a generic/universal way but basically we have a perf hit of having to write that data back to the host synchronously. How bad the hit is depends on the hardware configuration of the IPU chips and the host, the size of the model, and of course how big the tensor being read is. In most cases it will not be trivial. It's hard to give an exact estimate because it depends on so many factors. If the tensor is small then not much but a tensor you are more likely to watch (i.e a weight) would probably be in low double figure % hit. Depends on a lot of factors though.

I think the problem is that since it could be quite bad we'd rather users didn't get it by default, like if their canned models come with a large number of metric calls. Could it work in reverse on the IPU, users opt into the metrics with the warning about speed? We can offer this functionality as a debug feature, having said that, either via the anchors or a as a custom write back.

[SeanNaren]

That makes sense, I'm thinking this current logic:

By default we no-op all self.log calls, with a warning that log calls have been disabled on IPUs to reduce performance hits (even a link to the issue or something).
We introduce a parameter called enable_ipu_logging or something, that if passed to the trainer will try to log values using this anchor technique.

How does this sound? Also you mention it being tied to poptorch 2.1, which I assume will be out next month or so?

[tchaton]

Hey @davors72,

Any updates on this ticket ?

Best,
T.C

[davors72]

Hey @tchaton ,

Sorry, it wasn't added to 2.1. That's on me for deprioritising it after I realised you could use on_train_batch_end to get per-batch level logging. However we should still enable the lower level primative so it can be supported directly, I will reproritise it.