GEA training improvements: DDP batch size, sampler epoch, smooth head

[pjreddie]

Summary

Three small fixes that came up during GEA ecosystem segmentation finetuning:

• data_module: batch_size config now means global batch size — automatically divided by world_size in multi-GPU training so the config means the same thing regardless of GPU count
• lightning_module: Fix distributed sampler epoch shuffling — also call set_epoch on the sampler (not just batch_sampler) so data shuffling varies each epoch in DDP
• segmentation: Add smooth_sigma option to SegmentationHead — applies differentiable Gaussian blur to logits before loss/softmax (used in GEA smooth head experiments)

Test plan

• Used in GEA hyperparameter search experiments (p6_smooth_unfreeze, p8_smooth_train, p9_smooth_unfreeze)
• Multi-GPU training verified with correct per-GPU batch sizes

🤖 Generated with Claude Code

[favyen2]

Currently we often set batch_size based on available GPU memory. I don't think the existing option should be changed in behavior; if desired you could deprecate the existing one and add per_gpu_batch_size and global_batch_size options to replace it, and then it should raise error if neither is set or if both are set.

[favyen2]

The segmentation and train_dataloader.sampler changes look good to me but I think the batch_size behavior should either remain the same or deprecate it and add local_batch_size + global_batch_size options (and deprecated batch_size option should set the local batch size).

meant to comment not approve