# Wrap RaySGD over Tune trainable class directly

**URL:** <https://discuss.ray.io/t/wrap-raysgd-over-tune-trainable-class-directly/627>\
**Category:** Ray Tune\
**Created:** [January 27, 2021, 8:20pm UTC](https://discuss.ray.io/t/wrap-raysgd-over-tune-trainable-class-directly/627 "2021-01-27T20:20:28Z")\
**Posts on this page:** 14\
**Page:** 1

<div class="post-metadata">

**Author:** ![Gautham\_Krishnan](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/gautham_krishnan/32/344_2.png) [@Gautham\_Krishnan](https://discuss.ray.io/u/Gautham_Krishnan)\
**Post date:** [January 27, 2021, 8:20pm UTC](https://discuss.ray.io/t/wrap-raysgd-over-tune-trainable-class-directly/627/1 "2021-01-27T20:20:28Z")

</div>

I was going through the ray sgd documentation and came across the segment where the TorchOperator based class can be converted to a tune Trainable supported class calling the as\_Trainable method. Is it then possible to use the save\_checkpoint method as well? (ray version - 1.0.1). I would like to have access to a method that returns a checkpoint path rather than just the state dict, which is the case with the TorchOperator

---

<div class="post-metadata">

**Author:** ![rliaw](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/rliaw/32/24_2.png) [@rliaw](https://discuss.ray.io/u/rliaw)\
**Post date:** [January 27, 2021, 8:26pm UTC](https://discuss.ray.io/t/wrap-raysgd-over-tune-trainable-class-directly/627/2 "2021-01-27T20:26:41Z")

</div>

Hi @Gautham_Krishnan , the `trainable` returned by `as_trainable` is actually a class. So you can subclass that trainable to do whatever you want (including overriding save\_checkpoint).

Hope that helps!

---

<div class="post-metadata">

**Author:** ![Gautham\_Krishnan](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/gautham_krishnan/32/344_2.png) [@Gautham\_Krishnan](https://discuss.ray.io/u/Gautham_Krishnan)\
**Post date:** [January 28, 2021, 4:08am UTC](https://discuss.ray.io/t/wrap-raysgd-over-tune-trainable-class-directly/627/3 "2021-01-28T04:08:21Z")

</div>

Hey I tried that, think I’m going wrong somewhere. An example snippet would help. Thanks! @rliaw

---

<div class="post-metadata">

**Author:** ![rliaw](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/rliaw/32/24_2.png) [@rliaw](https://discuss.ray.io/u/rliaw)\
**Post date:** [January 28, 2021, 4:59am UTC](https://discuss.ray.io/t/wrap-raysgd-over-tune-trainable-class-directly/627/4 "2021-01-28T04:59:33Z")

</div>

Can you actually provide a bit more context about what you’re trying to do?

i.e., it’d probably be more effective if you showed me the code that currently doesn’t work (so that I understand your intent), than for me to show you code that might not actually be relevant 🙂

---

<div class="post-metadata">

**Author:** ![Gautham\_Krishnan](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/gautham_krishnan/32/344_2.png) [@Gautham\_Krishnan](https://discuss.ray.io/u/Gautham_Krishnan)\
**Post date:** [January 28, 2021, 5:44am UTC](https://discuss.ray.io/t/wrap-raysgd-over-tune-trainable-class-directly/627/5 "2021-01-28T05:44:53Z")

</div>

```
class TrainDistributed(TrainingOperator):
  def setup(self, config):
     <Some init code>

    # Register model, optimizer, and loss.
    self.model, self.optimizer, self.criterion = self.register(
        models=model,
        optimizers=optimizer,
        criterion=loss
    )

    # Register data loaders.
    self.register_data(train_loader=train_loader, validation_loader=val_loader)

  def train_batch(self, batch, batch_info):
    <over ride train_batch of training_operator>
    return packet

  def validate_batch(self, batch, batch_info):
    <over ride val_batch of training operator>
    return packet

  def save_checkpoint(self, checkpoint_dir):
    <I want to define this here and it should be recognized by tune. 
     Right now it doesn't happen and uses the standard state_dict
    provided by TorchTrainer. Namely, I can override state_dict, 
    but I would like to override this.>
    
    return checkpoint_dir
    
 def main():

  ray.init()

  TorchTrainable = TorchTrainer.as_trainable(
    training_operator_cls=TrainDistributed,
    use_gpu=True,
    num_workers=2,
    use_tqdm=True
 )
 analysis = tune.run(
    TorchTrainable,
    config=json.load(open('<path>')),
    stop={"training_iteration": 5},
    checkpoint_freq=2
)

```

Hope this helps. This saves the wts as required. But I want to save other files relevant to a checkpoint within the checkpoint folder, hence would like the path accessed rather than just the state dict

---

<div class="post-metadata">

**Author:** ![rliaw](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/rliaw/32/24_2.png) [@rliaw](https://discuss.ray.io/u/rliaw)\
**Post date:** [January 31, 2021, 9:16am UTC](https://discuss.ray.io/t/wrap-raysgd-over-tune-trainable-class-directly/627/6 "2021-01-31T09:16:50Z")

</div>

> [@Gautham\_Krishnan](#):
>
> ```auto
> provided by TorchTrainer. Namely, I can override state_dict, 
> 
> ```

Can you override `def state_dict` of the TrainingOperator and then inside the state\_dict function, do:

```auto
def state_dict(...):
    super_state = super().state_dict()
    my_file = write_my_file_to_a_tempdir()
    out = io.BytesIO()
    out.write(open(my_file).read())
    super_state["file_contents"] = out.getvalue()

```

---

<div class="post-metadata">

**Author:** ![Gautham\_Krishnan](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/gautham_krishnan/32/344_2.png) [@Gautham\_Krishnan](https://discuss.ray.io/u/Gautham_Krishnan)\
**Post date:** [February 1, 2021, 7:26am UTC](https://discuss.ray.io/t/wrap-raysgd-over-tune-trainable-class-directly/627/7 "2021-02-01T07:26:21Z")

</div>

Sure, will try this out. As a follow up, let’s say I already subclassed tune.trainable and using it for a single gpu training. Is there a way to directly plug in this class to TorchTrainer and enable Distributed training? I would like to add some code to def cleanup as well and I don’t see how that can be done with this flow @rliaw

---

<div class="post-metadata">

**Author:** ![rliaw](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/rliaw/32/24_2.png) [@rliaw](https://discuss.ray.io/u/rliaw)\
**Post date:** [February 1, 2021, 5:43pm UTC](https://discuss.ray.io/t/wrap-raysgd-over-tune-trainable-class-directly/627/8 "2021-02-01T17:43:19Z")

</div>

If you already have a training function (not trainable), and are using it for single gpu training, you can use Tune’s [Distributed Torch integrations instead](https://docs.ray.io/en/master/tune/api_docs/trainable.html#distributed-torch)

Here’s an example of how you might use this: [ddp\_mnist\_torch — Ray v2.0.0.dev0](https://docs.ray.io/en/master/tune/examples/ddp_mnist_torch.html)

---

<div class="post-metadata">

**Author:** ![Gautham\_Krishnan](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/gautham_krishnan/32/344_2.png) [@Gautham\_Krishnan](https://discuss.ray.io/u/Gautham_Krishnan)\
**Post date:** [February 1, 2021, 7:40pm UTC](https://discuss.ray.io/t/wrap-raysgd-over-tune-trainable-class-directly/627/9 "2021-02-01T19:40:45Z")

</div>

If I use the function approach, looks like I’ll have to call pytorch DDP on my own? Asking this looking at the example link^. As well as converting model to model.module? Would appreciate it if you can point me to the differences between using and not using ray-sgd for distributed training. Thanks @rliaw

---

<div class="post-metadata">

**Author:** ![rliaw](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/rliaw/32/24_2.png) [@rliaw](https://discuss.ray.io/u/rliaw)\
**Post date:** [February 1, 2021, 11:26pm UTC](https://discuss.ray.io/t/wrap-raysgd-over-tune-trainable-class-directly/627/10 "2021-02-01T23:26:35Z")

</div>

Yeah, you’ll have to call DDP on your own.

The main difference for using this vs raysgd is that you don’t have to restructure your standard training loop (while in torchtrainer you have to restructure your code into a training operator).

It also supports arbitrary checkpoints here, as you can see from the example. Hope that helps!

---

<div class="post-metadata">

**Author:** ![Gautham\_Krishnan](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/gautham_krishnan/32/344_2.png) [@Gautham\_Krishnan](https://discuss.ray.io/u/Gautham_Krishnan)\
**Post date:** [February 2, 2021, 4:07am UTC](https://discuss.ray.io/t/wrap-raysgd-over-tune-trainable-class-directly/627/11 "2021-02-02T04:07:12Z")

</div>

Okay got it. Would I have to handle combining the results of each worker in the function approach? Or is that handled internally?

---

<div class="post-metadata">

**Author:** ![rliaw](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/rliaw/32/24_2.png) [@rliaw](https://discuss.ray.io/u/rliaw)\
**Post date:** [February 2, 2021, 5:03am UTC](https://discuss.ray.io/t/wrap-raysgd-over-tune-trainable-class-directly/627/12 "2021-02-02T05:03:26Z")

</div>

I think only the rank 0 worker results are reported.

Optionally, if you want, you can use one of the torch aggregation primitives (like torch.distributed.reduce [Distributed communication package - torch.distributed — PyTorch 1.7.0 documentation](https://pytorch.org/docs/stable/distributed.html#backends)) to aggregate results.

---

<div class="post-metadata">

**Author:** ![Gautham\_Krishnan](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/gautham_krishnan/32/344_2.png) [@Gautham\_Krishnan](https://discuss.ray.io/u/Gautham_Krishnan)\
**Post date:** [February 8, 2021, 2:32pm UTC](https://discuss.ray.io/t/wrap-raysgd-over-tune-trainable-class-directly/627/13 "2021-02-08T14:32:07Z")

</div>

Thank you. I would have to call Distributed Data Sampler on my own too right, if I want to lock subsets of data to each worker? This has not been spoken about in the example I think?

---

<div class="post-metadata">

**Author:** ![rliaw](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/rliaw/32/24_2.png) [@rliaw](https://discuss.ray.io/u/rliaw)\
**Post date:** [February 9, 2021, 10:26pm UTC](https://discuss.ray.io/t/wrap-raysgd-over-tune-trainable-class-directly/627/14 "2021-02-09T22:26:10Z")

</div>

yeah, you would have to do that.