Matrix compression for federated learning broadcasts

src/main/java/org/apache/sysds/api/DMLScript.java

@@ -67,6 +67,7 @@
 import org.apache.sysds.runtime.controlprogram.context.SparkExecutionContext;
 import org.apache.sysds.runtime.controlprogram.federated.FederatedData;
 import org.apache.sysds.runtime.controlprogram.federated.FederatedWorker;
+import org.apache.sysds.runtime.controlprogram.federated.compression.CompressionType;
 import org.apache.sysds.runtime.controlprogram.federated.monitoring.FederatedMonitoringServer;
 import org.apache.sysds.runtime.controlprogram.federated.monitoring.models.CoordinatorModel;
 import org.apache.sysds.runtime.controlprogram.parfor.util.IDHandler;
@@ -131,6 +132,11 @@ public class DMLScript
 	public static double      GPU_MEMORY_UTILIZATION_FACTOR = 0.9;
 	// Set GPU memory allocator to use
 	public static String      GPU_MEMORY_ALLOCATOR       = "cuda";
+	// Enable/disable federated compression
+	public static boolean FEDERATED_COMPRESSION = false;
+	public static CompressionType FEDERATED_COMPRESSION_TYPE = CompressionType.TOPK;
+	public static double FEDERATED_COMPRESSION_SPARSITY = 0.01;
+	public static int FEDERATED_COMPRESSION_BITS = 4;
 	// Enable/disable lineage tracing
 	public static boolean     LINEAGE                    = DMLOptions.defaultOptions.lineage;
 	// Enable/disable deduplicate lineage items

src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedRequest.java

@@ -33,6 +33,7 @@
 import org.apache.sysds.runtime.controlprogram.caching.CacheBlock;
 import org.apache.sysds.runtime.controlprogram.caching.CacheDataOutput;
 import org.apache.sysds.runtime.controlprogram.caching.LazyWriteBuffer;
+import org.apache.sysds.runtime.controlprogram.federated.compression.CompressedMatrix;
 import org.apache.sysds.runtime.controlprogram.parfor.util.IDHandler;
 import org.apache.sysds.runtime.instructions.cp.ScalarObject;
 import org.apache.sysds.runtime.lineage.Lineage;
@@ -149,8 +150,8 @@ public long getChecksum(int i) {
 	}
 
 	private void calcChecksum() throws IOException {
-		for (Object ob : _data) {
-			if (!(ob instanceof CacheBlock) && !(ob instanceof ScalarObject))
+		for(Object ob : _data) {
+			if(!(ob instanceof CacheBlock) && !(ob instanceof ScalarObject) && !(ob instanceof CompressedMatrix))
 				continue;
 
 			Checksum checksum = new Adler32();
@@ -174,6 +175,21 @@ private void calcChecksum() throws IOException {
 					throw new IOException("Failed to serialize cache block.", ex);
 				}
 			}
+
+			if(ob instanceof CompressedMatrix) {
+				try {
+					java.io.ByteArrayOutputStream baos = new java.io.ByteArrayOutputStream();
+					java.io.ObjectOutputStream oos = new java.io.ObjectOutputStream(baos);
+					oos.writeObject(ob);
+					oos.flush();
+					byte[] bytes = baos.toByteArray();
+					checksum.update(bytes, 0, bytes.length);
+					_checksums.add(checksum.getValue());
+				}
+				catch(Exception ex) {
+					throw new IOException("Failed to serialize CompressedMatrix.", ex);
+				}
+			}
 		}
 	}
 
@@ -187,6 +203,8 @@ public long estimateSerializationBufferSize() {
 			for(Object obj : _data) {
 				if(obj instanceof CacheBlock)
 					minBufferSize += ((CacheBlock<?>)obj).getExactSerializedSize();
+				else if(obj instanceof CompressedMatrix)
+					minBufferSize += 1024 * 1024; // conservative 1MB estimate for compressed matrix
 			}
 		}
 		if(_lineageTrace != null)

src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederatedWorkerHandler.java

@@ -50,6 +50,8 @@
 import org.apache.sysds.runtime.controlprogram.context.SparkExecutionContext;
 import org.apache.sysds.runtime.controlprogram.federated.FederatedRequest.RequestType;
 import org.apache.sysds.runtime.controlprogram.federated.FederatedResponse.ResponseType;
+import org.apache.sysds.runtime.controlprogram.federated.compression.CompressedMatrix;
+import org.apache.sysds.runtime.controlprogram.federated.compression.CompressionFactory;
 import org.apache.sysds.runtime.controlprogram.federated.monitoring.models.DataObjectModel;
 import org.apache.sysds.runtime.controlprogram.federated.monitoring.models.EventModel;
 import org.apache.sysds.runtime.controlprogram.federated.monitoring.models.EventStageModel;
@@ -488,7 +490,21 @@ private FederatedResponse putVariable(FederatedRequest request, ExecutionContext
 		// wrap transferred cache block into cacheable data
 		Data data;
 		long size = 0;
-		if(v instanceof CacheBlock) {
+		if(v instanceof CompressedMatrix) {
+			try {
+				CompressedMatrix cm = (CompressedMatrix) v;
+				org.apache.sysds.runtime.controlprogram.federated.compression.MatrixCompressor compressor = CompressionFactory
+					.createForDecompression(cm.getType());
+				org.apache.sysds.runtime.matrix.data.MatrixBlock decompressed = compressor.decompress(cm);
+				var block = ExecutionContext.createCacheableData(decompressed);
+				size = block.getDataSize();
+				data = block;
+			}
+			catch(Exception e) {
+				throw new FederatedWorkerHandlerException("Failed to decompress matrix: " + e.getMessage());
+			}
+		}
+		else if(v instanceof CacheBlock) {
 			var block = ExecutionContext.createCacheableData((CacheBlock<?>) v);
 			size = block.getDataSize();
 			data = block;

src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederationMap.java

@@ -42,12 +42,18 @@
 import org.apache.sysds.hops.fedplanner.FTypes.AlignType;
 import org.apache.sysds.hops.fedplanner.FTypes.FType;
 import org.apache.sysds.lops.RightIndex;
+import org.apache.sysds.api.DMLScript;
 import org.apache.sysds.runtime.DMLRuntimeException;
 import org.apache.sysds.runtime.controlprogram.caching.CacheBlock;
 import org.apache.sysds.runtime.controlprogram.caching.CacheableData;
 import org.apache.sysds.runtime.controlprogram.federated.FederatedRequest.RequestType;
 import org.apache.sysds.runtime.instructions.cp.CPOperand;
 import org.apache.sysds.runtime.instructions.cp.ScalarObject;
+import org.apache.sysds.runtime.controlprogram.federated.compression.CompressionConfig;
+import org.apache.sysds.runtime.controlprogram.federated.compression.CompressionFactory;
+import org.apache.sysds.runtime.controlprogram.federated.compression.CompressedMatrix;
+import org.apache.sysds.runtime.controlprogram.federated.compression.MatrixCompressor;
+import org.apache.sysds.runtime.matrix.data.MatrixBlock;
 import org.apache.sysds.runtime.instructions.cp.VariableCPInstruction;
 import org.apache.sysds.runtime.lineage.LineageItem;
 import org.apache.sysds.runtime.util.IndexRange;
@@ -138,6 +144,25 @@ private FederatedRequest broadcast(CacheableData<?> data, LineageItem lineageIte
 		// is fine, because with broadcast all data on all workers)
 		data.setFedMapping(copyWithNewIDAndRange(
 			cb.getNumRows(), cb.getNumColumns(), id, FType.BROADCAST));
+
+		// === COMPRESSION INTEGRATION ===
+		// Compress matrix before network transfer; worker decompresses on receipt
+		if(DMLScript.FEDERATED_COMPRESSION && cb instanceof MatrixBlock) {
+			try {
+				CompressionConfig config = CompressionConfig.builder().enable(true)
+					.withType(DMLScript.FEDERATED_COMPRESSION_TYPE)
+					.withSparsity(DMLScript.FEDERATED_COMPRESSION_SPARSITY)
+					.withBits(DMLScript.FEDERATED_COMPRESSION_BITS).build();
+				MatrixCompressor compressor = CompressionFactory.create(config);
+				CompressedMatrix compressed = compressor.compress((MatrixBlock) cb);
+				return new FederatedRequest(RequestType.PUT_VAR, lineageItem, id, compressed);
+			}
+			catch(Exception ex) {
+				// Fall back to uncompressed on any error
+			}
+		}
+		// === END COMPRESSION INTEGRATION ===
+
 		return new FederatedRequest(RequestType.PUT_VAR, lineageItem, id, cb);
 	}
 

src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressedMatrix.java

@@ -0,0 +1,96 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.sysds.runtime.controlprogram.federated.compression;
+
+import java.io.Serializable;
+
+/**
+ * Generic container for compressed matrix data. Stores the compressed representation along with metadata needed for
+ * decompression and size estimation.
+ */
+public class CompressedMatrix implements Serializable {
+
+	private static final long serialVersionUID = 1L;
+
+	private final CompressionType type;
+	private final int numRows;
+	private final int numCols;
+	private final Object compressedData; // Technique-specific data
+	private final double compressionRatio;
+	private final byte[] metadata; // Optional: scaling factors, etc.
+
+	public CompressedMatrix(CompressionType type, int numRows, int numCols, Object compressedData,
+		double compressionRatio) {
+		this(type, numRows, numCols, compressedData, compressionRatio, null);
+	}
+
+	public CompressedMatrix(CompressionType type, int numRows, int numCols, Object compressedData,
+		double compressionRatio, byte[] metadata) {
+		this.type = type;
+		this.numRows = numRows;
+		this.numCols = numCols;
+		this.compressedData = compressedData;
+		this.compressionRatio = compressionRatio;
+		this.metadata = metadata;
+	}
+
+	public CompressionType getType() {
+		return type;
+	}
+
+	public int getNumRows() {
+		return numRows;
+	}
+
+	public int getNumCols() {
+		return numCols;
+	}
+
+	public Object getCompressedData() {
+		return compressedData;
+	}
+
+	public double getCompressionRatio() {
+		return compressionRatio;
+	}
+
+	public byte[] getMetadata() {
+		return metadata;
+	}
+
+	/** Estimate original size in bytes (8 bytes per double) */
+	public long estimateOriginalSizeBytes() {
+		return (long) numRows * numCols * 8;
+	}
+
+	/** Estimate compressed size in bytes */
+	public long getCompressedSizeBytes() {
+		if(compressedData instanceof byte[]) {
+			return ((byte[]) compressedData).length;
+		}
+		return 0;
+	}
+
+	@Override
+	public String toString() {
+		return String.format("CompressedMatrix[%s, %dx%d, ratio=%.2fx]", type.getId(), numRows, numCols,
+			compressionRatio);
+	}
+}

src/main/java/org/apache/sysds/runtime/controlprogram/federated/compression/CompressionConfig.java

@@ -0,0 +1,113 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.sysds.runtime.controlprogram.federated.compression;
+
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Immutable configuration for compression in federated operations. Uses the Builder pattern for flexible, readable
+ * configuration.
+ *
+ * Usage example: CompressionConfig config = CompressionConfig.builder() .enable(true) .withType(CompressionType.TOPK)
+ * .withSparsity(0.01) .build();
+ */
+public class CompressionConfig {
+
+	private final boolean enabled;
+	private final CompressionType type;
+	private final Map<String, Object> parameters;
+
+	private CompressionConfig(Builder builder) {
+		this.enabled = builder.enabled;
+		this.type = builder.enabled ? builder.type : CompressionType.NONE;
+		this.parameters = new HashMap<>(builder.parameters);
+	}
+
+	public boolean isEnabled() {
+		return enabled;
+	}
+
+	public CompressionType getType() {
+		return type;
+	}
+
+	public Map<String, Object> getParameters() {
+		return new HashMap<>(parameters);
+	}
+
+	/** Convenience getter for sparsity parameter (TopK) */
+	public double getSparsity() {
+		return (double) parameters.getOrDefault("sparsity", 0.01);
+	}
+
+	/** Convenience getter for bits parameter (Quantization) */
+	public int getBits() {
+		return (int) parameters.getOrDefault("bits", 4);
+	}
+
+	@Override
+	public String toString() {
+		return String.format("CompressionConfig[enabled=%s, type=%s, params=%s]", enabled, type.getId(), parameters);
+	}
+
+	// -----------------------------------------------------------------------
+	// Builder
+	// -----------------------------------------------------------------------
+
+	public static Builder builder() {
+		return new Builder();
+	}
+
+	public static class Builder {
+		private boolean enabled = false;
+		private CompressionType type = CompressionType.NONE;
+		private final Map<String, Object> parameters = new HashMap<>();
+
+		public Builder enable(boolean enabled) {
+			this.enabled = enabled;
+			return this;
+		}
+
+		public Builder withType(CompressionType type) {
+			this.type = type;
+			return this;
+		}
+
+		public Builder withParameter(String key, Object value) {
+			this.parameters.put(key, value);
+			return this;
+		}
+
+		/** Shorthand for TopK sparsity ratio */
+		public Builder withSparsity(double sparsity) {
+			return withParameter("sparsity", sparsity);
+		}
+
+		/** Shorthand for quantization bit width */
+		public Builder withBits(int bits) {
+			return withParameter("bits", bits);
+		}
+
+		public CompressionConfig build() {
+			return new CompressionConfig(this);
+		}
+	}
+}