Explicitly allow or disallow agents with costs in different frameworks

[nicola-bastianello]

At this time we don't explicitly allow or disallaw different agents to have costs in different frameworks:

• In some places we raise when the frameworks are mismatched (like when using Cost._validate_cost_operation(self, other, check_framework: bool = True)). This would probably disallow mismatched frameworks when e.g. computing the sum of the local costs.
• In other places, we implicitly assume that all costs have the same framework, like when we set SumCost.framework as return self.costs[0].framework. This doesn't explicitly disallow mismatched frameworks, but silently ignores it

Overall, I think we should decide for a clear behavior and implement it everywhere: either we allow or we disallow mismatched frameworks, and ensure that the docs clarify this.

Thanks to the iop layer and the Array object, I think it's not impossible to allow mismatched frameworks, but there are some places like SumCost where it's not clear what we should do.

[Simpag]

Speaking more broadly and not just about individual agents:

Its generally not recommended to use a combination of different frameworks because its much slower. Some frameworks also don’t like to be used together, for example tensorflow and jax both pre-allocate gpu memory which in turn crashes the program because they together try to use 2x amount of available memory.

However, we could leave it and allow for it as most modules are built to handle multiple frameworks. This would be less limiting in terms of combining cost functions. However, I would imagine that combining cost functions will only be used for small scale simple experimentation and once a good cost has been found users would implement it directly as one cost function, maybe using a gpu accelerated framework. I wouldn't recommend using the cost functions to apply regularization to neural networks as it is much more efficient to do that natively using the loss functions.

[Simpag]

The interoperability layer does not support the use of two different frameworks in one call. For example, iop.stack cannot be called with a list of values from different frameworks. Personally I would disallow the use of different frameworks but still allow for different cost functions on each agent.

[nicola-bastianello]

The interoperability layer does not support the use of two different frameworks in one call. For example, iop.stack cannot be called with a list of values from different frameworks. Personally I would disallow the use of different frameworks but still allow for different cost functions on each agent.

ok, I think this is a good compromise if we want to allow working with mixed cost functions. but then we need to decide how to handle SumCost.framework, e.g. converting to a single framework? but that probably defeats the purpose of using costs from different frameworks

[Simpag]

ok, I think this is a good compromise if we want to allow working with mixed cost functions. but then we need to decide how to handle SumCost.framework, e.g. converting to a single framework? but that probably defeats the purpose of using costs from different frameworks

I cant really see a benefit of using cost functions from different frameworks. It will just slow down the algorithms too much. I would probably mention in the docs that the feature of combining cost functions is only meant to be used with simple cost functions, more advanced costs, such as PyTorchCost, should not be combined.

[nicola-bastianello]

one (very ambitious) thing we could do: switch to cost functions entirely based on iop, and when combining them they use the same framework/device. e.g.: we create a PyTorchCost and add a iop-based regularizer to it: it will return a regularizer with framework=pytorch and same device. or better yet, it modifies the loss function of PyTorchCost directly

[nicola-bastianello]

if we do iop-based costs, we could then have something like set_seed which globlly sets framework and device. this way the can be removed as args from iop functions

[Simpag]

I thought that the regularizers already were framework agnostic? So adding a regularizer to a PyTorch cost will make it use PyTorch operations but it is still much slower than just using a loss function with the regularizer. I believe that almost all PyTorch loss functions allow for regularization, and if they don't it is easy to create custom losses and add it yourself.

I don't think we can make every cost using iop; or, we could but it would come at the cost of performance as we cannot optimize it as easily and much higher code complexity. For example in PyTorchCost: we would not be able to use dataloaders, won't be able to easily do the optimized per-sample gradient calculations, no easy way to create and initialize a model for the cost function. We would need to create iop functions for most linear algebra functions and solvers, we would need to expose auto-gradient functionality which would have very limited support. I don't think this is a very feasible solution, I think the performance hit and the reduction of applications (can't use framework specific optimisations and utilities) will limit the framework more than the fact that you can't combine more niche cost functions.

Regarding removing framework and device args from iop functions, the idea I originally had is that users should use as many _like functions as possible so they don't have to worry about framework and device args. I didn't want to expose any functions which had the framework and device parameters, but at some point it was needed. I'm not sure if it is anymore.

[nicola-bastianello]

I thought that the regularizers already were framework agnostic? So adding a regularizer to a PyTorch cost will make it use PyTorch operations but it is still much slower than just using a loss function with the regularizer. I believe that almost all PyTorch loss functions allow for regularization, and if they don't it is easy to create custom losses and add it yourself.

ah yes indeed, my mistake

I don't think we can make every cost using iop; or, we could but it would come at the cost of performance as we cannot optimize it as easily and much higher code complexity. For example in PyTorchCost: we would not be able to use dataloaders, won't be able to easily do the optimized per-sample gradient calculations, no easy way to create and initialize a model for the cost function. We would need to create iop functions for most linear algebra functions and solvers, we would need to expose auto-gradient functionality which would have very limited support. I don't think this is a very feasible solution, I think the performance hit and the reduction of applications (can't use framework specific optimisations and utilities) will limit the framework more than the fact that you can't combine more niche cost functions.

yes that's a good point, it doesn't make sense to sacrifice performance so much.

Regarding removing framework and device args from iop functions, the idea I originally had is that users should use as many _like functions as possible so they don't have to worry about framework and device args. I didn't want to expose any functions which had the framework and device parameters, but at some point it was needed. I'm not sure if it is anymore.

I think it is still needed (e.g. when creating the initial conditions for agents states). if we could pass a cost in the _like functions we might be able to do everything with _like, but I'm not sure it's a very elegant solution

[Simpag]

Regarding removing framework and device args from iop functions, the idea I originally had is that users should use as many _like functions as possible so they don't have to worry about framework and device args. I didn't want to expose any functions which had the framework and device parameters, but at some point it was needed. I'm not sure if it is anymore.

I think it is still needed (e.g. when creating the initial conditions for agents states). if we could pass a cost in the _like functions we might be able to do everything with _like, but I'm not sure it's a very elegant solution

In the backend it is definitely needed to keep the code clean, but I was thinking more about the user-exposed API. But it might still be needed in some algorithms. I don't think its needed in cost functions as users can just use framework-specific functions

[adrianardv]

Okey, If I have understood correctly you've talked about 3 different things:

1) Networks with agents using different frameworks

Case where agent1.cost.framework != agent2.cost.framework
Right now, this is allowed in construction, though network tensors like weights and adjacency are created in the first agent's framework. Then algorithms fail later when they stack/mix states/messages.

Places where compatibility is assumed by reading metadata from the first agent/client instead of validating all agents:

• networks.py: in P2PNetwork weights and adjacency; in FedNetwork default server initialization.

So, if we are not going to support mixed framework network object, maybe we should raise an error earlier, at construction? If this is the case, then:

• Add a network level validation step in networks.py that checks every agent has the same cost.shape, cost.framework, and cost.device, then the existing "use first agent" helpers can be kept.
• Keep the current "use first agent" logic, becuase after validation it becomes safe.

If we were to allow it, then:

• send/ _send_one would need receiver-side conversion.
• Fed algorithm aggregation stacks (iop.stack) client updates and therefore assumes one backend. ( As Simon said, iop.stack cannot be called with a list of values from different frameworks.)
• Mean and consensus metrics also stack all agetn's x

2) Agents in a network with different cost functions (same framework/device):

The codebase right now already allows different agents to carry different Cost objects and different concrete cost classes. Practical requirement: all agents optimize the same parameter shape, and share framework/device.

3) Composition of costs from different frameworks:

If you build a single cost out of multiple sub-costs, are those sub-costs allowed to come from different array frameworks? This is where the code is inconsistent right now.

• SumCost "allows" the sum of costs of different frameworks. However, SumCost silently reports the framework/device is the first term's framework. Also, in SumCost.gradient/hessian(), sub-results are aggregated with iop.stack, which does not support stacking mixed framework arrays in one call.
• In LinearRegressionCost/LogisticRegressionCost generic addition same framework/device is not enforced consistently.
• for a _CompositeRegularizerCost to be formed, same framework/device is already enforced (raise error).
• EmpiricalRegularizedCost, cost and regularizer need to have the same framework/device. It is already enforced (raise error)

If we explicitly disallow 3):

• Disallow framework/device mismatches inside a composed cost, SumCost. --> This can be done changing _validate_cost_operation() so framework and device are checked by default. And explicitly validating in SumCost.__init__ that all terms have the same shape, framework and device.
• Keep framework and device as "first term" because that will now be guaranteed.

About the iop-based costs:

• I agree with simon on keeping backend-bound costs specialized.
• Regularizers are already the closest thing to framework-agnostic costs in the codebase, since their methods are iop-based. But they still have explicit framework/device metadata.
• So if a regularizer is created with matching framework/device, composing it with another cost is already supported by the current design, although for backend-bound costs such as PyTorchCost it is still not the most efficient path compared with native framework-specific regularization.

---

So, my understanding is that the simplest consistent rule is to enforce backend/device homogeneity both at the network level and at the composed-cost level:

• Disallow mixed-framework/device networks (1)
• Disallow mixed-framework/device composed costs (3)
• Keep different cost classes per agent supported, as long as shape/framework/device match (2)

[Simpag]

I agree with @adrianardv. We should disallow mixing of frameworks and devices throughout. There is no real reason to allow for this since everything can be done using the same framework and at best allowing it only reduces performance, while it might not even work in some cases. Tensorflow is always compiled during runtime so I am not even sure how well it would work with mixing frameworks or customizing models. Allowing for different cost classes for each agent should be completely fine as long as they use the same framework/device.

[nicola-bastianello]

thanks @adrianardv for the clear summary!

I agree with you both: let's disallow mixing in networks and disallow mixing in composite costs

for iop: it would still be nice to have something like set_seed for framework+device. this could be called automatically when costs are generated, and future checks refer to these as the correct framework+device. so we don't ask the users to directly specify framework and device unless they need to