Integrate distributed inference with chat/server #1381

mreso · 2024-11-16T05:45:06Z

This PR integrates distributed inference with the CLI options chat and server and implements Option 2b of #1376.

To test run on machine with 4 gpus:

python  torchchat.py chat llama3.1  --max-new-tokens 10  --distributed --tp 2 --pp 2
python  torchchat.py server llama3.1   --distributed --tp 2 --pp 2

pytorch-bot · 2024-11-16T05:45:09Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchchat/1381

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 1 Pending, 1 Unrelated Failure

As of commit 7650153 with merge base 90749d2 ():

NEW FAILURE - The following job has failed:

Run the aoti runner with CUDA using stories / test-runner-aot-cuda / linux-job (gh)
RuntimeError: Command docker exec -t 0ed4d44989c6d6b67cf6d67209ac615036fb2b370a98e10b23877b199007da18 /exec failed with exit code 1

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

pull / runner-aoti (16-core-ubuntu) (gh) (trunk failure)
Process completed with exit code 1.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

…action

mikekgfb · 2024-11-18T03:58:12Z

torchchat/cli/builder.py

+            "Meta-Llama-3-70B": "meta-llama/Meta-Llama-3-70B-Instruct",
+            "Meta-Llama-3.1-70B": "meta-llama/Meta-Llama-3.1-70B-Instruct",
+
+        }


What's the gap to using the models described in model_config/models.json. (as implied by TODO comment)

Definitely should not be part of the present PR, but I think as a north star, it would be desirable to grab the same models (and download and mgmt infra etc) for non-distributed and distributed.

Thanks for the comment @mikekgfb. The gap wasn't that big, the args.model just wasn't accessible at that point and I wanted to take a deeper look to fix it right. Removed the cruel hack and I now save the distribution_path when creating the builder_arg. Still not sure if this is the intended way of locating the checkpoint though.

Not sure what you're trying to do. If it's just that the dict/json data structure describing the mapping isn't in scope, maybe what you want is some methods that give you the relevant info?

Also, these seem to be mapping of short names to HF network paths - should we not have a way to [ick them from the local filesystem (since the torchchat cli already manages download and all that). Oh, and if the answer is "we have bigger fish to fry, hooking this up is not highest priority" I will wholeheartedly agree. This is more about understanding the context of this PR.

Where I'm lacking the context is how you go from all the weights being available locally on a node to reading those weights on another node? And maybe that's why you prefer to straight up pick the files from HF? (Although local distribution from an already downloaded set of weights probs has higher bandwidth?) Again, there's much bigger fish to fry, and I think this PR is a good step in the direction of frying those fish ;)

Oh I see, let me try to fill in some context for this PR. Previously, the distributed inference solution lived completely separate in its own script (within the torchchat repo but) separate from torchchat.py. Distributed provides its own utils to load either hf or torchchat weights (where the torchchat part is currently broken IIRC). In a previous PR (#1327), I enabled the usage of torchchat.py generate with a distributed model. This PR only progresses the integration into the cli by enabling chat/server but stops short from replacing the weight loading mechanics which are still custom to distributed.

So, yes, for this PR I was only looking for a quick and dirty way to map arg.model_name -> "huggingface distribution str" (without actually having model_name at hand) to load the weights from the hf cache. I now modified the PR to use the information provided in the model_config/models.json as you suggested. In a next PR we should then alter torchchat/distributed/checkpoint_utils.py to leverage torchchat infra like (e.g. builder_args) to locate and access the files.

Jack-Khuu

Partial review so far, still need to dig through generate

cc: @Gasoonjia as secondary reviewer

torchchat/usages/openai_api.py

torchchat/usages/server.py

…action

torchchat/generate.py

Jack-Khuu · 2024-12-11T02:58:35Z

torchchat/generate.py

@@ -358,7 +368,6 @@ def prefill(
        sequential_prefill=True,
        **sampling_kwargs,
    ) -> torch.Tensor:
-        # logging.debug(f"x: {x}, input_pos: {input_pos}")


Heads up the logging changes might get merge conflicted with #1336

torchchat/generate.py

Jack-Khuu · 2024-12-11T18:00:53Z

torchchat/generate.py

+        if self.pp_rank == self.last_pp_rank:
+            new_token, next_prob = self.sample(logits, need_probs=need_probs, **sampling_kwargs)
+        else:
+            new_token = torch.zeros(1, 1, device=self.device, dtype=torch.int64)
+
+        if self.pp_rank == self.last_pp_rank and self.pp_rank != self.first_pp_rank:
+            dist.send(
+                new_token,
+                dst=self.first_pp_rank_global_id,
+                group=self.pp_group,
+            )
+        elif self.pp_rank == self.first_pp_rank and self.pp_rank != self.last_pp_rank:
+            dist.recv(
+                new_token,
+                src=self.last_pp_rank_global_id,
+                group=self.pp_group,
+            )
+            #TODO: Why do we get 2d tensor here?
+            new_token=new_token[0]


Ditto above as prefill

torchchat/generate.py

Jack-Khuu · 2024-12-11T19:28:42Z

torchchat/generate.py

+    else:
+        run_generator(args)


Shuold setup_env also be triggering the generator?

Maybe it's just a naming change (run_with_setup~?), but setup_env doesn't make me think that the fn is called

Yes, valid point. I've renamed it to run_in_dist_env to reflect that it runs the target as well as connecting it to distributed.

Jack-Khuu · 2024-12-11T19:31:02Z

torchchat/generate.py

-from torchchat.distributed.generate import DistributedGenerator
+from torchchat.distributed.utils import (
+    Color as color,
+    setup_env,


nit: setup_env is pretty generic, let's alias it so it's obviously for a distributed context

Jack-Khuu · 2024-12-11T19:32:02Z

torchchat/usages/server.py

+    if builder_args.distributed:
+        world_size = builder_args.tp * builder_args.pp
+        mp_context = mp.get_context('spawn')
+        queue = mp_context.Queue()
+
+    if builder_args.distributed:


Identical conditional check

…inference_without_abstraction

Co-authored-by: Jack-Khuu <jack.khuu.7@gmail.com>

…n the comments as well

…the target

…action

mikekgfb · 2024-12-20T00:06:34Z

This PR integrates distributed inference with the CLI options chat and server and implements Option 2b of #1376.

To test run on machine with 4 gpus:
python  torchchat.py chat llama3.1  --max-new-tokens 10  --distributed --tp 2 --pp 2
python  torchchat.py server  llama3.1  --max-new-tokens 10  --distributed --tp 2 --pp 2

Should torchchat generate work as well, code would suggest that we can do like so (not sure about max new tokens?):

python  torchchat.py generate llama3.1 --distributed --tp 2 --pp 2

mreso · 2024-12-20T00:25:51Z

@mikekgfb Yes, distributed generate was already working before this PR, thats why I did not list it separately.

mreso added 2 commits November 15, 2024 18:04

Integrate distributed inference without introducing abstraction

7b4d5c5

Cleanup old distributed inference integration

e7670c3

facebook-github-bot added the CLA Signed This label is managed by the Meta Open Source bot. label Nov 16, 2024

mreso requested review from lessw2020 and Jack-Khuu November 16, 2024 05:45

Merge branch 'main' into refactor/distributed_inference_without_abstr…

08a8e03

…action

mreso mentioned this pull request Nov 16, 2024

[WIP] Generate base class for better integration of distributed inference #1355

Closed

mikekgfb reviewed Nov 18, 2024

View reviewed changes

mreso added 2 commits November 18, 2024 10:49

Read distribution from model_config

d5bca9b

Declare distribution_path if args.model is not given

76895cc

Jack-Khuu reviewed Nov 19, 2024

View reviewed changes

torchchat/usages/openai_api.py Outdated Show resolved Hide resolved

torchchat/usages/openai_api.py Outdated Show resolved Hide resolved

torchchat/usages/server.py Show resolved Hide resolved

torchchat/usages/server.py Outdated Show resolved Hide resolved

mreso and others added 4 commits November 19, 2024 14:50

Address some nits from PR review

3ef1296

Merge branch 'main' into refactor/distributed_inference_without_abstr…

04cdfd0

…action

Merge branch 'main' into refactor/distributed_inference_without_abstr…

924a096

…action

Merge branch 'main' into refactor/distributed_inference_without_abstr…

99c33e8

…action

Jack-Khuu reviewed Dec 11, 2024

View reviewed changes

mreso and others added 11 commits December 16, 2024 10:56

Merge remote-tracking branch 'origin/main' into refactor/distributed_…

773894f

…inference_without_abstraction

Added comment on model size all reduce + type hint

7cb98c9

Apply suggestions from code review

10fb55a

Co-authored-by: Jack-Khuu <jack.khuu.7@gmail.com>

Make sure speculative decoding is disable for pp >1 and remark this i…

28d7836

…n the comments as well

Refactor conditions in pp

68eec0b

Rename and alter signature of setup_env to reflect that it also runs …

3ad31e8

…the target

Rename setup_env in server + fix condition

e07b03d

Merge branch 'main' into refactor/distributed_inference_without_abstr…

daf902c

…action

Merge branch 'main' into refactor/distributed_inference_without_abstr…

db5fd1b

…action

Update generate.py

7ac16f9

Add default value to add_generation_prompt to preserve bc

7650153

Jack-Khuu approved these changes Dec 19, 2024

View reviewed changes

Jack-Khuu merged commit cc0ffce into pytorch:main Dec 19, 2024
51 of 53 checks passed

mikekgfb mentioned this pull request Dec 19, 2024

[RFC] Integration of Distributed Inference into TorchChat #1376

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Integrate distributed inference with chat/server #1381

Integrate distributed inference with chat/server #1381

mreso commented Nov 16, 2024 •

edited

Loading

pytorch-bot bot commented Nov 16, 2024 •

edited

Loading

mikekgfb Nov 18, 2024 •

edited

Loading

mreso Nov 18, 2024

mikekgfb Nov 19, 2024

mreso Nov 20, 2024

Jack-Khuu left a comment

Jack-Khuu Dec 11, 2024

Jack-Khuu Dec 11, 2024

Jack-Khuu Dec 11, 2024

mreso Dec 17, 2024

Jack-Khuu Dec 11, 2024

Jack-Khuu Dec 11, 2024

mikekgfb commented Dec 20, 2024 •

edited

Loading

mreso commented Dec 20, 2024

Integrate distributed inference with chat/server #1381

Integrate distributed inference with chat/server #1381

Conversation

mreso commented Nov 16, 2024 • edited Loading

pytorch-bot bot commented Nov 16, 2024 • edited Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchchat/1381

❌ 1 New Failure, 1 Pending, 1 Unrelated Failure

mikekgfb Nov 18, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Jack-Khuu left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mikekgfb commented Dec 20, 2024 • edited Loading

mreso commented Dec 20, 2024

mreso commented Nov 16, 2024 •

edited

Loading

pytorch-bot bot commented Nov 16, 2024 •

edited

Loading

mikekgfb Nov 18, 2024 •

edited

Loading

mikekgfb commented Dec 20, 2024 •

edited

Loading