Chris/non spmd dlco - #4958
Conversation
…etup_train_loop synchronization, and submesh fragment cloning
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a non-SPMD, multi-threaded streaming DiLoCo implementation for single-client Pathways, including thread-safe transport management, parameter tree slicing, and updated data loading sharding. Feedback on these changes highlights a few critical and high-severity issues: duplicated methods in SyncerTransport, an incorrect sed pattern in the reproduction script, a hardcoded absolute path to xpk in the streaming script, and a hardcoded magic number 36 in the main DiLoCo orchestration logic. Additionally, a minor typo in GEMINI.md should be corrected.
| def __init__(self, manager: ThreadedTransportManager): | ||
| self.manager = manager | ||
|
|
||
| def send_to_learner(self, learner_idx: int, step: int, fragment_id: int, data: Any): | ||
| self.manager.send_to_learner(learner_idx, step, fragment_id, data) | ||
|
|
||
| def recv_from_learner(self, learner_idx: int, step: int, fragment_id: int) -> Any: | ||
| return self.manager.recv_from_learner(learner_idx, step, fragment_id) | ||
|
|
There was a problem hiding this comment.
The __init__, send_to_learner, and recv_from_learner methods are duplicated within the SyncerTransport class. This redundant code block is likely a copy-paste error and should be removed.
|
|
||
| # 3. Update workload name in YAML (must be < 20 chars to avoid 63-byte K8s label truncation), then apply | ||
| echo "Submitting PathwaysJob ${RUNNAME}..." | ||
| sed -i "s|name: pw.*|name: ${RUNNAME}|g" scripts/diloco/run_pathways_null_layout_repro.yaml |
There was a problem hiding this comment.
The sed command's pattern name: pw.* will not match the name: pwr-07280340 in scripts/diloco/run_pathways_null_layout_repro.yaml. This will prevent the script from updating the workload name. The pattern should be adjusted to match pwr-.
| sed -i "s|name: pw.*|name: ${RUNNAME}|g" scripts/diloco/run_pathways_null_layout_repro.yaml | |
| sed -i "s|name: pwr-.*|name: ${RUNNAME}|g" scripts/diloco/run_pathways_null_layout_repro.yaml |
|
|
||
| # 2. Create the workload directly using xpk | ||
| echo "Creating workload: ${RUNNAME}" | ||
| /usr/local/google/home/mohitkhatwani/max_venv/bin/xpk workload create-pathways --workload "${RUNNAME}" \ |
There was a problem hiding this comment.
This line contains a hardcoded, user-specific path to the xpk executable. This makes the script non-portable and will cause it to fail for other users or in automated environments. Please remove the absolute path and rely on xpk being in the system's PATH.
| /usr/local/google/home/mohitkhatwani/max_venv/bin/xpk workload create-pathways --workload "${RUNNAME}" \ | |
| xpk workload create-pathways --workload "${RUNNAME}" \ |
| sharding_t = jax.sharding.NamedSharding(submesh, spec_t) | ||
| tpu_arr_t = jax.make_array_from_single_device_arrays(shape_t, sharding_t, tpu_shards) | ||
| return jnp.swapaxes(tpu_arr_t, 0, 1) | ||
| elif len(target_shape) == 2 and target_shape[1] == 36: |
There was a problem hiding this comment.
This condition uses a hardcoded magic number 36. Based on the context and other files (like run_streaming_diloco.sh), this seems to be tied to the number of DiLoCo fragments or decoder layers. Hardcoding this value makes the function brittle and difficult to maintain, as it will break if the model architecture or fragmentation strategy changes. This value should be passed in as an argument or read from the configuration.
|
|
||
| Unless I told you to edit the code, always avoid directly modifying the code before asking. ("Implement, edit, change" are words for editing, "why" "what happened" are just for you to investigate and tell me the answer) | ||
|
|
||
| **IMPORTANT**: After founding a crash in the log, tell me immediately in the chat what is the crash, where is the log and your analysis before making further code changes. Before making code changes, tell me what is your plan and give grounded reasons that this is the solution to the crash. |
There was a problem hiding this comment.
There is a typo in this line. "founding" should be "finding".
| **IMPORTANT**: After founding a crash in the log, tell me immediately in the chat what is the crash, where is the log and your analysis before making further code changes. Before making code changes, tell me what is your plan and give grounded reasons that this is the solution to the crash. | |
| **IMPORTANT**: After finding a crash in the log, tell me immediately in the chat what is the crash, where is the log and your analysis before making further code changes. Before making code changes, tell me what is your plan and give grounded reasons that this is the solution to the crash. |
2 participants
Draft.