Can gold_trainer support privileged context?

[zouYC2021]

Current gold_trainer(trl/experimental/gold/gold_trainer.py) seems to only support distillation between students and teachers that have the same input. I wonder whether it can support teacher model with privileged context as in sdft_trainer.

[matdou]

It's not supported right now. The GOLD trainer assumes symmetric inputs. Supporting privileged context would mean allowing separate teacher/student inputs in the batch and running the two forward passes independently before computing KL, similar to what sdft_trainer does. Would be a reasonable addition.

Not a trivial change (data collator + forward pass split) but well-scoped with sdft_trainer as reference. Happy to take a stab at it if there's interest.

[zouYC2021]

It's not supported right now. The GOLD trainer assumes symmetric inputs. Supporting privileged context would mean allowing separate teacher/student inputs in the batch and running the two forward passes independently before computing KL, similar to what sdft_trainer does. Would be a reasonable addition.

Not a trivial change (data collator + forward pass split) but well-scoped with sdft_trainer as reference. Happy to take a stab at it if there's interest.

That would be quite nice if someone can try that.