[gss10282023]

Summary

With DDP, the gathered contrastive loss only backpropagates through local embeddings, so after DDP's averaging its gradient is 1/world_size of the single-GPU gradient, while the RTD loss is not affected. On 2 GPUs, lambda_weight=0.005 acts like 0.01.

This PR multiplies the contrastive loss by dist.get_world_size() during distributed training, before lambda_weight * RTD is added. Only the contrastive term is multiplied, not the combined loss. Single-GPU training is unchanged.

The PR also reads five training metrics from model.module when CLTrainer receives a DDP wrapper. The original reads raise AttributeError before the first distributed update, so this small prerequisite lets the corrected loss path run through the native Trainer. On a single GPU the reads still use model.

Issue links

Fixes #29

Validation

At 33b29a38, one AdamW step with a global batch of 64 on 1 vs. 2 GPUs, lambda_weight=0.005, the same number of RTD tokens per GPU: before the change the full gradient differs from the 1-GPU run by about 0.49999; after the change the difference is 1.563e-7 in FP32 and 1.989e-16 in FP64, and the update and AdamW state match as well. Reproduced on a second machine.

The recorded experiments used the equivalent five trainer-read adaptation now included in this PR. They also sized electra_head to the small test model and ran without --batchnorm; those two experimental controls are not part of this PR. This exact combined PR diff has not had a fresh full Trainer run.

Checks on this branch

• Both changed files parse, and the five metric reads resolve on a plain model and on a wrapper exposing .module in a small check.
• The numerical values above come from earlier runs of the equivalent contrastive multiplier and metric adaptation. I have not run the full combined PR diff on this branch.

Prepared with AI assistance; opened as a draft for code review.