Add GPU & Distributed Training Support

[JoeHolt]

No description provided.

[JoeHolt]

I believe the only relevant file from my other PR was the _dist.py, so I moved it over here into this cleaner PR

[JoeHolt]

When we last met, you asked that I update the API of DaskClassifier to look like this:

DaskClassifier(..., batch_size="geodamp", batch_size__latency=60, batch_size__factor=5)

I am have a few questions:

• How should I implement the latency? Should I start tracking the current epoch internally within the class? Ie every time run_single_epoch is called, increment internal epoch count by 1.
• How can I implement the specific dampening algorithm? Ie what the easiest way to make _dist.py use GeoDamp vs PadaDamp?

[stsievert]

relevant file from my other PR

Does tests/test_dist_damping.py need to be moved over from #15 too?

How should I implement the latency delay?

Something like this:

def initialize(self, X_train):
    self._meta = {..., "num_examples": 0}
def _partial_fit(self, X, y):
    self._meta["num_examples"] += len(X)
    _epochs = self._meta["num_examples"] / len(X_train)
    ...

Ie what the easiest way to make _dist.py use GeoDamp vs PadaDamp?

I think the easiest way is:

class GeoDamp:
    def __init__(self, delay=60, factor=5, initial=128):
        self.delay = delay
        self.factor = factor
        self.initial = initial

    def damp(self, meta: Dict[str, Any]) -> int:
        # meta is DaskBaseDamper._meta
        epochs = meta["epochs"]
        return self.initial * (self.factor ** (epochs // self.delay))

This will require some parsing of the arguments. DaskClassifier(batch_size="geodamp", batch_size__delay=70) should turn into DaskClassifier(batch_size=GeoDamp(delay=70)).

[JoeHolt]

I implemented the new API based on your description but had some questions:

• It seems like it makes more sense to set the batch size arguments during a call to partial_fit vs fit. Does this make sense?
• How can I adapt the classes in damping.py to work with dist.py? those classes seem like they need a lot more init arguments than the simple example we used.
• Do I need to then update the test files?

[stsievert]

I'd put the dampers (BaseDamper, Damper, ...) in a new file, _dampers.py.

The current tests should pass. I'd make a new test to ensure that GeoDamp works as expected. I think test_damping.py has some code that will be useful.

[stsievert]

Suggested change

	grads_per_worker=32,
	grads_per_worker: int = 128,

[stsievert]

This should be set in initialize.

The Scikit-learn API says that parameters ending in underscores should only be set when fit is called (https://scikit-learn.org/stable/glossary.html#term-attributes)

[stsievert]

Suggested change

	def get_batch_size(self, batch_size, kwargs: Dict[str, Any]) -> BaseDamper:
	def _get_damper(self, batch_size, kwargs: Dict[str, Any]) -> BaseDamper:

[stsievert]

This chunk of code should go in initialize.

[stsievert]

Suggested change

	def run_single_epoch(self, X, y=None, **fit_params):
	def _run_single_epoch(self, X, y=None, **fit_params):

[stsievert]

I see – there's already BaseDamper and GeoDamp. It might be okay to remove this from __init__.py. Users can always do from adadamp.dampers import GeoDamp.

[stsievert]

Suggested change

	from adadamp.adadamp.dampers import SimpleBaseDamper, SimpleGeoDamp
	from .dampers import SimpleBaseDamper, SimpleGeoDamp

[stsievert]

It helps to run pip install -e . in the root of this repo.

[JoeHolt]

What does this do?

[stsievert]

pip install -e /path/to/adadamp installs the AdaDamp package. It looks like from adadamp.adadamp.dampers import ... is not using the adadamp package; instead, it looks like it's using a relative path.

It might cleaner to do from adadamp.dampers import SimpleBaseDamper, SimpleGeoDamp. That's what from .dampers import is doing (but is relative, not absolute). Here's more detail: https://realpython.com/absolute-vs-relative-python-imports/

[stsievert]

Suggested change:

+ if not isinstance(self.batch_size, (str, int, np.integer, SimpleBaseDamper)):
+     raise ValueError("self.batch_size needs to be ...")
...
+ elif isinstance(self.batch_size, (int, np.integer)):
...
- if not isinstance(self.batch_size_, SimpleBaseDamper):
-     raise ValueError("self.batch_size is not ...")