Move local weights construction to drive_dfmm