Skip to content

Update tf pin to 03/2023#4840

Merged
JackCaoG merged 25 commits intomasterfrom
tf-update-032023
Apr 20, 2023
Merged

Update tf pin to 03/2023#4840
JackCaoG merged 25 commits intomasterfrom
tf-update-032023

Conversation

@wonjoo-wj
Copy link
Copy Markdown
Collaborator

@wonjoo-wj wonjoo-wj commented Mar 30, 2023 •
edited
Loading

Update tf pin to 03/2023

@wonjoo-wj wonjoo-wj force-pushed the tf-update-032023 branch 4 times, most recently from 8d0a541 to f1e37bd Compare April 4, 2023 21:53
@lsy323 lsy323 force-pushed the tf-update-032023 branch from f1e37bd to 9053703 Compare April 4, 2023 23:33
@will-cromar
Copy link
Copy Markdown
Collaborator

will-cromar commented Apr 6, 2023

Can you update the libtpu version too?

xla/setup.py

Line 62 in 29a8bf2

_libtpu_version = '0.1.dev20230213'

@wonjoo-wj
Copy link
Copy Markdown
Collaborator Author

wonjoo-wj commented Apr 6, 2023

Yep, we'll update that.

@wonjoo-wj wonjoo-wj force-pushed the tf-update-032023 branch 2 times, most recently from 1e10aac to 68a61d1 Compare April 10, 2023 23:05
@wonjoo-wj wonjoo-wj self-assigned this Apr 14, 2023
@wonjoo-wj wonjoo-wj force-pushed the tf-update-032023 branch 2 times, most recently from 9391930 to 717d9a3 Compare April 17, 2023 09:13
@JackCaoG
Copy link
Copy Markdown
Collaborator

JackCaoG commented Apr 17, 2023

@wonjoolee95 Can you rebase this pr?

@will-cromar will-cromar mentioned this pull request Apr 17, 2023
Merged
@wonjoo-wj wonjoo-wj force-pushed the tf-update-032023 branch 2 times, most recently from d5aa01e to e3965e2 Compare April 19, 2023 06:27
JackCaoG
JackCaoG reviewed Apr 19, 2023
View reviewed changes
Comment thread .circleci/common.sh
PJRT_DEVICE=GPU python test/test_train_mp_imagenet_fsdp.py --fake_data --auto_wrap_policy type_based --use_small_fake_sample --num_epochs=1
XLA_DISABLE_FUNCTIONALIZATION=1 PJRT_DEVICE=GPU python test/test_train_mp_imagenet_fsdp.py --fake_data --use_nested_fsdp --use_small_fake_sample --num_epochs=1
# This test fails on GPU with 03/30 TF-pin update (https://github.com/pytorch/xla/pull/4840)
# XLA_DISABLE_FUNCTIONALIZATION=1 PJRT_DEVICE=GPU python test/test_train_mp_imagenet_fsdp.py --fake_data --use_nested_fsdp --use_small_fake_sample --num_epochs=1
Copy link
Copy Markdown
Collaborator

@JackCaoG JackCaoG Apr 19, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TODO: reword the comment to say this is only happening on CI, can't repo on GPU VM.(do that in a separate pr)

Also do we really need this test, seems like the only difference is to disable the functionization. @alanwaketan

JackCaoG
JackCaoG reviewed Apr 19, 2023
View reviewed changes
Comment thread test/test_torch_distributed_fsdp_meta.py
JackCaoG
JackCaoG reviewed Apr 19, 2023
View reviewed changes
Comment thread tf_patches/local_rendezvous.diff
"This may indicate a bug on the caller side. (b/274683676)";
}
- return rc_keep_alive;
+ return std::make_optional<tsl::core::RefCountPtr<Rendezvous>>(std::move(rc_keep_alive));
Copy link
Copy Markdown
Collaborator

@JackCaoG JackCaoG Apr 19, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this should be fixed in tf right? so we can remove this patch in next release.

Copy link
Copy Markdown
Collaborator Author

@wonjoo-wj wonjoo-wj Apr 19, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, hopefully we can also remove the triton patch in the next tf update.

JackCaoG
JackCaoG reviewed Apr 19, 2023
View reviewed changes
Comment thread tf_patches/topk_rewriter.diff
};
- return Match(sort->operand(1), match_iota(data->shape().dimensions())) ||
- Match(sort->operand(1), m::Broadcast(match_iota(sort_dims)));
+ return Match(sort->operand(1), match_iota(absl::Span<const int64_t>(data->shape().dimensions().begin(), data->shape().dimensions().size()))) ||
Copy link
Copy Markdown
Collaborator

@JackCaoG JackCaoG Apr 19, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ditto, should be fixed in tf

JackCaoG
JackCaoG reviewed Apr 19, 2023
View reviewed changes
Comment thread third_party/xla_client/pjrt_computation_client.cc
// TODO(yeounoh) currently only support single-slice execution
execute_options.multi_slice_config = nullptr;

// Required as of cl/518733871
Copy link
Copy Markdown
Collaborator

@JackCaoG JackCaoG Apr 19, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do not include cl number in open pr, use pr number on public github instead.(fix in next pr).

JackCaoG
JackCaoG reviewed Apr 19, 2023
View reviewed changes
Comment thread third_party/xla_client/pjrt_computation_client.cc
execute_options.untuple_result = options.explode_tuple;
execute_options.strict_shape_checking = false;

// Required as of cl/518733871
Copy link
Copy Markdown
Collaborator

@JackCaoG JackCaoG Apr 19, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same as below

JackCaoG
JackCaoG approved these changes Apr 19, 2023
View reviewed changes
Copy link
Copy Markdown
Collaborator

@JackCaoG JackCaoG left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Mostly lgtm, you can fix comments in a separate pr

@wonjoo-wj
Copy link
Copy Markdown
Collaborator Author

wonjoo-wj commented Apr 19, 2023

Thanks for the comments, @JackCaoG.

Adding the log for the failing GPU tests (test/test_torch_distributed_fsdp_meta.py):

2023-04-11 02:48:11.603946: F  256355 external/org_tensorflow/tensorflow/compiler/xla/service/layout_assignment.cc:157] Check failed: ShapeUtil::Compatible(shape_layout.shape(), instruction->operand(operand_no)->shape()) f32[128]{0} is not compatible with f32[256]{0} (for operand 0 of instruction %reduce-scatter.180 = f32[128]{0} reduce-scatter(f32[256]{0} %add.175), replica_groups={}, constrain_layout=true, dimensions={0}, to_apply=%AddComputation.29)
2023-04-11 02:48:11.604912: F  256356 external/org_tensorflow/tensorflow/compiler/xla/service/layout_assignment.cc:157] Check failed: ShapeUtil::Compatible(shape_layout.shape(), instruction->operand(operand_no)->shape()) f32[128]{0} is not compatible with f32[256]{0} (for operand 0 of instruction %reduce-scatter.180 = f32[128]{0} reduce-scatter(f32[256]{0} %add.175), replica_groups={}, constrain_layout=true, dimensions={0}, to_apply=%AddComputation.29)

@wonjoo-wj
Copy link
Copy Markdown
Collaborator Author

wonjoo-wj commented Apr 19, 2023

Tests on GPU python test/test_train_mp_imagenet_fsdp.py --fake_data --use_nested_fsdp --use_small_fake_sample --num_epochs=1 failing with same reason.. This is a bit time wasting, as I can't run reproduce on my own GPUVM docker, pushed a new commit to disable this for now.

@JackCaoG JackCaoG marked this pull request as ready for review April 20, 2023 00:53
@wonjoo-wj wonjoo-wj changed the title [WIP] Update tf pin to 03/2023 Update tf pin to 03/2023 Apr 20, 2023
@wonjoo-wj
Copy link
Copy Markdown
Collaborator Author

wonjoo-wj commented Apr 20, 2023 •
edited
Loading

I'll merge this tomorrow morning, as all CIs are green. Merged.

Thanks for the review and comments, Jack. I'll also work with the next tf-pin rotation to make sure we try to remove some of these pathces.

@JackCaoG JackCaoG merged commit c1138e8 into master Apr 20, 2023
This was referenced May 26, 2023
Closed
Open
@lsy323 lsy323 deleted the tf-update-032023 branch March 4, 2024 19:12
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@JackCaoG JackCaoG JackCaoG approved these changes

Assignees

@wonjoo-wj wonjoo-wj

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@wonjoo-wj @will-cromar @JackCaoG